From 72fe1b623f72e0116b04f203f286e05472606741 Mon Sep 17 00:00:00 2001 From: heyuanliu-intel Date: Tue, 25 Aug 2026 16:06:17 +0000 Subject: [PATCH] [Bugfix][Hardware][Intel] MiniMax-H3: reuse the tiled-VAE gather buffer across requests on XPU ## Purpose **Problem:** the tiled video VAE gathers every rank's decoded tiles with a variable-shape all-gather whose output is allocated afresh each call. After the per-request `empty_cache()` the next gather lands at a new device address. On XPU the collective backend registers a device-side resource per new receive address and does not reclaim it, so device memory outside `torch.xpu.memory_allocated()` grows every request until the card is full. (Mechanism inferred from intervention experiments, not observed directly.) **Fix:** keep one landing buffer per (shape, dtype, device) on the adapter and gather into it with `all_gather_into_tensor`; tiles are cloned out because the buffer is reused. Installed only when the adapter's device is XPU, bound on the checkpoint instance; other backends keep the checkpoint's own method. The buffer cache is uncapped on purpose: the key is the tile geometry, and evicting would re-register a new address. ## Test Plan **Reproduce:** 4 XPU cards, TP4/SP1, model-level offload on, send three identical t2va requests (1344x768, 209 frames, 2 steps) and sample each rank's non-PyTorch device memory between requests: it climbs monotonically while `memory_allocated()` stays flat. ```bash pytest -q tests/diffusion/models/minimax_h3/test_minimax_h3_vae_gather_reuse.py ``` **vLLM Version:** as pinned by `requirements/`. **vLLM-Omni Commit:** rebased onto current `main`. ## Test Result CPU tests pass (XPU-only install, tile order for even and uneven shares, buffer reuse at same geometry, tiles survive the next gather). On the 4-card A/B, non-PyTorch peak growth from request 1 to 2 dropped from 6617.69 MiB to 4008.76 MiB per rank (-39.4%); worst-rank free memory at request 2 rose from 2.93 MiB to 1662.50 MiB. Growth is reduced, not eliminated. Signed-off-by: heyuanliu-intel --- .../test_minimax_h3_vae_gather_reuse.py | 376 ++++++++++++++++++ vllm_omni/diffusion/models/minimax_h3/vae.py | 123 ++++++ 2 files changed, 499 insertions(+) create mode 100644 tests/diffusion/models/minimax_h3/test_minimax_h3_vae_gather_reuse.py diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_gather_reuse.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_gather_reuse.py new file mode 100644 index 00000000000..160d7903e51 --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_gather_reuse.py @@ -0,0 +1,376 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""CPU tests for the pinned tiled-VAE gather buffer in the MiniMax-H3 video VAE. + +The checkpoint's ``_all_gather_tiled_results`` allocates its gather output on +every call, so with a per-request ``empty_cache()`` the collective lands on a +new device address every request and XCCL's per-address registrations pile up. +On XPU the adapter replaces that method with an equal-shape gather into a +buffer it keeps alive. These tests pin the properties that make it a fix +rather than a rewrite: the address is stable, the tiles come back +bit-identical and in the same order, the caller owns what it gets back, and no +other device is touched. +""" + +import logging +import sys +import types +from typing import Any + +import pytest +import torch +import torch.distributed as dist + +from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module +from vllm_omni.diffusion.models.minimax_h3.vae import MiniMaxH3VideoVAE + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + +TILE = (1, 3, 2, 4, 4) + + +class _CheckpointGatherSentinel: + """Stands in for the checkpoint's own bound method. + + Identity is the point: a test asserts this exact object is still installed + on the devices the adapter must leave alone. + """ + + def __call__(self, tasks, num_tiles): # pragma: no cover - never invoked + raise AssertionError("the checkpoint gather must not run in these tests") + + +class _FakeCheckpointModel: + def __init__(self): + self._all_gather_tiled_results = _CheckpointGatherSentinel() + + +def _local_tiles(all_tiles, num_tiles, rank, sp_size): + """One rank's round-robin share, as ``klvae.py:275`` computes it.""" + return [all_tiles[index] for index in range(rank, num_tiles, sp_size)] + + +def _reference_gather(all_tiles, num_tiles, sp_size): + """The checkpoint's variable-shape list gather, reimplemented verbatim. + + Written from the control flow of ``_all_gather_tiled_results`` in the + checkpoint's ``klvae.py`` (lines 252-273) -- stack each rank's share, gather + the list, walk it back into global tile order -- so that it is an + independent oracle rather than a restatement of the code under test. + """ + gathered = [torch.stack(_local_tiles(all_tiles, num_tiles, rank, sp_size), dim=0) for rank in range(sp_size)] + results = [None] * num_tiles + for rank, rank_tensors in enumerate(gathered): + for k in range(rank_tensors.shape[0]): + global_index = k * sp_size + rank + if global_index >= num_tiles: + break + results[global_index] = rank_tensors[k] + return results + + +def _peer_stacks(all_tiles, num_tiles, sp_size): + """What every rank hands to the collective, padded the way the adapter pads.""" + max_tasks = -(-num_tiles // sp_size) + stacks = [] + for rank in range(sp_size): + share = _local_tiles(all_tiles, num_tiles, rank, sp_size) + stacked = share[0].new_zeros((max_tasks, *share[0].shape)) + stacked[: len(share)] = torch.stack(share, dim=0) + stacks.append(stacked) + return stacks + + +def _tiles(num_tiles, seed=0): + generator = torch.Generator().manual_seed(seed) + return [torch.rand(TILE, generator=generator) for _ in range(num_tiles)] + + +@pytest.fixture +def fake_collective(monkeypatch): + """Replace the collective with a single-process stand-in. + + ``all_gather_into_tensor`` is the only distributed call on this path, so + filling the output buffer from a caller-supplied peer list exercises the + real padding, keying, addressing and unpacking code without a process + group. + """ + calls: list[dict[str, Any]] = [] + + def fake_all_gather_into_tensor(output, input_tensor, group=None, async_op=False): + peers = calls[-1]["peers"] + assert output.shape == (len(peers), *input_tensor.shape) + calls[-1]["sent"] = input_tensor.clone() + calls[-1]["buf_ptr"] = output.data_ptr() + for rank, peer in enumerate(peers): + output[rank].copy_(peer) + + monkeypatch.setattr(dist, "all_gather_into_tensor", fake_all_gather_into_tensor) + return calls + + +class _Records(logging.Handler): + """vllm's logger sets ``propagate=False``, so caplog never sees these.""" + + def __init__(self): + super().__init__(level=logging.DEBUG) + self.messages = [] + self.records = [] + + def emit(self, record): + self.messages.append(record.getMessage()) + self.records.append((record.levelno, record.getMessage())) + + +@pytest.fixture +def gather_log(): + handler = _Records() + logger = vae_module.logger + previous = logger.level + logger.addHandler(handler) + logger.setLevel(logging.DEBUG) + try: + yield handler + finally: + logger.removeHandler(handler) + logger.setLevel(previous) + + +def _vae(sp_size, sp_rank=0): + """An adapter instance with the override installed, no checkpoint needed.""" + state = { + "sp_size": sp_size, + "sp_rank": sp_rank, + "sp_enabled": sp_size > 1, + "sp_process_group": object(), + } + module = types.ModuleType("fake_gather_ckpt.parallel") + module.get_parallel_state = lambda: state # type: ignore[attr-defined] + sys.modules.setdefault("fake_gather_ckpt", types.ModuleType("fake_gather_ckpt")) + sys.modules["fake_gather_ckpt.parallel"] = module + + vae = object.__new__(MiniMaxH3VideoVAE) + vae.remote = type("Remote", (), {"__module__": "fake_gather_ckpt.klvae"})() + vae.model = _FakeCheckpointModel() + vae._device_target = torch.device("xpu") + vae._tile_gather_buffers = {} + vae._tile_gather_stats = {"hits": 0, "allocs": 0, "bytes": 0} + vae._checkpoint_tile_gather = None + vae._install_persistent_tile_gather() + return vae + + +def _run(vae, calls, all_tiles, num_tiles, sp_size, sp_rank): + calls.append({"peers": _peer_stacks(all_tiles, num_tiles, sp_size)}) + return vae.model._all_gather_tiled_results(_local_tiles(all_tiles, num_tiles, sp_rank, sp_size), num_tiles) + + +# -------------------------------------------------------------------------- +# The address is pinned across calls. This is what the accumulation is about: +# drop the dictionary lookup and every call registers a new receive address. +# -------------------------------------------------------------------------- + + +def test_the_landing_buffer_address_is_stable_across_calls(fake_collective): + vae = _vae(4) + pointers = [] + for request in range(3): + _run(vae, fake_collective, _tiles(8, seed=request), 8, 4, 0) + pointers.append(fake_collective[-1]["buf_ptr"]) + + assert len(set(pointers)) == 1, f"gather buffer moved across calls: {pointers}" + buffer = next(iter(vae._tile_gather_buffers.values())) + assert vae._tile_gather_stats == { + "hits": 2, + "allocs": 1, + "bytes": buffer.numel() * buffer.element_size(), + } + + +def test_each_distinct_shape_gets_its_own_pinned_buffer(fake_collective): + """One decode gathers several temporal clips, and the head/tail clips differ.""" + vae = _vae(4) + _run(vae, fake_collective, _tiles(8), 8, 4, 0) + first = fake_collective[-1]["buf_ptr"] + _run(vae, fake_collective, _tiles(12), 12, 4, 0) + second = fake_collective[-1]["buf_ptr"] + _run(vae, fake_collective, _tiles(8, seed=5), 8, 4, 0) + third = fake_collective[-1]["buf_ptr"] + + assert first != second + assert third == first + assert vae._tile_gather_stats["allocs"] == 2 + assert vae._tile_gather_stats["hits"] == 1 + + +# -------------------------------------------------------------------------- +# Equivalence with the checkpoint's own gather, padded branch included. A +# reordered tile grid is silent at runtime, so order is asserted as well. +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize("num_tiles", [5, 8, 9, 28]) +@pytest.mark.parametrize("sp_rank", [0, 3]) +def test_gathered_tiles_match_the_checkpoint_gather(fake_collective, num_tiles, sp_rank): + sp_size = 4 + all_tiles = _tiles(num_tiles, seed=num_tiles) + vae = _vae(sp_size, sp_rank=sp_rank) + + produced = _run(vae, fake_collective, all_tiles, num_tiles, sp_size, sp_rank) + expected = _reference_gather(all_tiles, num_tiles, sp_size) + + assert len(produced) == num_tiles + assert all(tile is not None for tile in produced) + for index, (got, want) in enumerate(zip(produced, expected)): + assert torch.equal(got, want), f"tile {index} differs from the checkpoint gather" + # ...and it is the tile that belongs at that index, not a permutation that + # happens to agree with an equally permuted oracle. + for index, tile in enumerate(produced): + assert torch.equal(tile, all_tiles[index]) + + +def test_the_padded_branch_is_actually_exercised(fake_collective): + """9 tiles over 4 ranks: ranks 1-3 are one short, so the padding runs.""" + sp_size, num_tiles = 4, 9 + all_tiles = _tiles(num_tiles, seed=9) + assert len(_local_tiles(all_tiles, num_tiles, 1, sp_size)) < -(-num_tiles // sp_size) + + vae = _vae(sp_size, sp_rank=1) + produced = _run(vae, fake_collective, all_tiles, num_tiles, sp_size, 1) + + sent = fake_collective[-1]["sent"] + assert sent.shape[0] == 3 + assert torch.equal(sent[2], torch.zeros_like(sent[2])) + for index, tile in enumerate(produced): + assert torch.equal(tile, all_tiles[index]) + + +def test_an_empty_local_share_still_raises(fake_collective): + """A rank with no tiles fails loudly, exactly as the checkpoint does.""" + vae = _vae(4) + with pytest.raises(ValueError, match="empty tasks on sp rank"): + vae.model._all_gather_tiled_results([], 8) + + +def test_more_tiles_than_round_robin_allows_is_refused(fake_collective): + """The unpacking arithmetic is only valid for the checkpoint's ownership.""" + vae = _vae(4) + with pytest.raises(ValueError, match="round-robin"): + vae.model._all_gather_tiled_results(_tiles(5), 8) + + +# -------------------------------------------------------------------------- +# Ownership: the buffer is overwritten by the next gather, so returned tiles +# must not alias it. +# -------------------------------------------------------------------------- + + +def test_returned_tiles_do_not_alias_the_pinned_buffer(fake_collective): + sp_size, num_tiles = 4, 8 + first_tiles = _tiles(num_tiles, seed=1) + vae = _vae(sp_size) + produced = _run(vae, fake_collective, first_tiles, num_tiles, sp_size, 0) + + buffer = next(iter(vae._tile_gather_buffers.values())) + buffer_storage = buffer.untyped_storage().data_ptr() + for tile in produced: + assert tile.untyped_storage().data_ptr() != buffer_storage + + # The next gather reuses that memory; the first call's tiles must survive it. + _run(vae, fake_collective, _tiles(num_tiles, seed=2), num_tiles, sp_size, 0) + for index, tile in enumerate(produced): + assert torch.equal(tile, first_tiles[index]) + + +# -------------------------------------------------------------------------- +# Non-XPU behaviour is unchanged: the checkpoint's own method stays installed. +# -------------------------------------------------------------------------- + + +def _construct(monkeypatch, device, model): + """Run the real ``__init__`` against a stub checkpoint component.""" + + class _Remote(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.zeros(1)) + self.model = model + + monkeypatch.setattr( + vae_module, + "_load_component_config", + lambda path: {"latent_channels": 1, "latents_mean": [0.0], "latents_std": [1.0]}, + ) + # **kwargs: the loader gained a trust_remote_code argument at some point, + # and this stub does not care either way. + monkeypatch.setattr(vae_module, "_load_remote_component", lambda path, config, **kwargs: _Remote()) + # The adapter only stores this handle; resolving it for real needs an + # accelerator, which a CPU test runner does not have. + monkeypatch.setattr(torch, "get_device_module", lambda *args, **kwargs: types.SimpleNamespace()) + return MiniMaxH3VideoVAE("unused", device=torch.device(device)) + + +@pytest.mark.parametrize("device", ["cpu", "meta"]) +def test_non_xpu_devices_keep_the_checkpoint_gather(monkeypatch, device): + model = _FakeCheckpointModel() + original = model._all_gather_tiled_results + + vae = _construct(monkeypatch, device, model) + + assert vae._tile_gather_reuse_enabled() is False + assert vae.model._all_gather_tiled_results is original + assert vae._checkpoint_tile_gather is None + + +@pytest.mark.parametrize( + ("device", "enabled"), + [("xpu", True), ("cpu", False), ("cuda", False), ("meta", False)], +) +def test_only_xpu_asks_for_the_pinned_buffer(device, enabled): + vae = object.__new__(MiniMaxH3VideoVAE) + vae._device_target = torch.device(device) + assert vae._tile_gather_reuse_enabled() is enabled + + +def test_construction_installs_the_override_when_the_device_asks_for_it(monkeypatch): + """Close the loop: ``__init__`` must consult the predicate, not ignore it.""" + monkeypatch.setattr(MiniMaxH3VideoVAE, "_tile_gather_reuse_enabled", lambda self: True) + model = _FakeCheckpointModel() + original = model._all_gather_tiled_results + + vae = _construct(monkeypatch, "cpu", model) + + assert vae.model._all_gather_tiled_results is not original + assert vae._checkpoint_tile_gather is original + + +# -------------------------------------------------------------------------- +# The decision-point receipt carries quantities, not just presence: a line +# that only says "installed" cannot tell reuse from reallocation. +# -------------------------------------------------------------------------- + + +def test_the_receipt_line_reports_hits_allocs_and_the_address(fake_collective, gather_log): + vae = _vae(4) + install = [ + (level, line) for level, line in gather_log.records if "persistent tile gather installed device=xpu" in line + ] + # The one-shot install line is what an operator needs at info level. + assert install and all(level == logging.INFO for level, _ in install) + gather_log.messages.clear() + gather_log.records.clear() + + _run(vae, fake_collective, _tiles(8), 8, 4, 0) + first_ptr = fake_collective[-1]["buf_ptr"] + _run(vae, fake_collective, _tiles(8, seed=3), 8, 4, 0) + + records = [(level, line) for level, line in gather_log.records if "[H3_VAE_GATHER] reuse=" in line] + # The per-gather line fires once per decoder tile batch, so it is debug. + assert all(level == logging.DEBUG for level, _ in records) + lines = [line for _, line in records] + assert len(lines) == 2 + assert "reuse=alloc" in lines[0] and "hits=0 allocs=1" in lines[0] + assert "reuse=hit" in lines[1] and "hits=1 allocs=1" in lines[1] + assert f"buf_ptr=0x{first_ptr:x}" in lines[0] + assert f"buf_ptr=0x{first_ptr:x}" in lines[1] diff --git a/vllm_omni/diffusion/models/minimax_h3/vae.py b/vllm_omni/diffusion/models/minimax_h3/vae.py index 2e830fa0481..56d8eea18b2 100644 --- a/vllm_omni/diffusion/models/minimax_h3/vae.py +++ b/vllm_omni/diffusion/models/minimax_h3/vae.py @@ -153,6 +153,11 @@ def __init__( self.use_slicing = False self.parallel_size = 1 self.device_module = torch.get_device_module() + self._tile_gather_buffers: dict[tuple[Any, ...], torch.Tensor] = {} + self._tile_gather_stats = {"hits": 0, "allocs": 0, "bytes": 0} + self._checkpoint_tile_gather = None + if self._tile_gather_reuse_enabled(): + self._install_persistent_tile_gather() def load_to_device(self) -> None: if self._stager is not None: @@ -218,6 +223,124 @@ def _native_parallel_state(self) -> dict[str, Any]: parallel_module = importlib.import_module(f"{package}.parallel") return parallel_module.get_parallel_state() + def _tile_gather_reuse_enabled(self) -> bool: + """Whether this device needs the tiled-VAE gather buffer pinned. + + XPU only. The accumulation this guards against is an XCCL-side + registration that is kept for every distinct receive-buffer address and + never reclaimed; no other backend in tree does that, so everywhere else + the checkpoint's own method is left in place and behaviour is unchanged. + """ + + return self._device_target.type == "xpu" + + def _install_persistent_tile_gather(self) -> None: + """Route the checkpoint's tiled-VAE gather through a pinned buffer. + + The checkpoint's ``_all_gather_tiled_results`` allocates its gather + output afresh on every call. Model-level offload calls ``empty_cache()`` + once per request, so the next request's output lands on a new device + address; XCCL registers a non-reclaimable resource per new receive + address, and the registrations accumulate until the card is full. The + replacement below keeps one landing buffer per (shape, dtype, device) + alive on this adapter, so the address the collective writes to is the + same one every request. + + The override is bound on the checkpoint *instance*, not its class: the + class is remote code shared by every component loaded from the same + checkpoint, and only this adapter knows the device it runs on. + """ + + self._checkpoint_tile_gather = self.model._all_gather_tiled_results + self.model._all_gather_tiled_results = self._persistent_tile_gather + logger.info( + "[H3_VAE_GATHER] persistent tile gather installed device=%s", + self._device_target.type, + ) + + def _persistent_tile_gather( + self, + tasks: list[torch.Tensor], + num_tiles: int, + ) -> list[torch.Tensor]: + """Equal-shape replacement for the checkpoint's tiled-result gather. + + Contract kept identical to the checkpoint's method: return a list of + ``num_tiles`` tensors in global tile order, and raise on an empty local + share so a rank that owns no tile cannot silently skip the collective. + + Tile ownership is round-robin (``range(sp_rank, num_tiles, sp_size)``), + so every rank can compute every other rank's task count from + ``num_tiles`` and ``sp_size`` alone. That makes the per-rank payloads + equal once the leading task dimension is padded to ``max_tasks``, which + is what lets a single ``all_gather_into_tensor`` into a pinned buffer + replace the variable-shape gather. + + Returned tiles are cloned out of the buffer: the buffer is overwritten + by the next call and callers hold the tiles past that point. + """ + + state = self._native_parallel_state() + group = state["sp_process_group"] + sp_size = int(state["sp_size"]) + sp_rank = int(state["sp_rank"]) + + if not tasks: + raise ValueError(f"Found empty tasks on sp rank {sp_rank}") + + max_tasks = -(-num_tiles // sp_size) + if len(tasks) > max_tasks: + raise ValueError( + f"sp rank {sp_rank} holds {len(tasks)} tiles but round-robin " + f"ownership of {num_tiles} tiles across {sp_size} ranks allows " + f"at most {max_tasks}" + ) + if len(tasks) == max_tasks: + stacked = torch.stack(tasks, dim=0) + else: + # Pad the leading (task) dimension only. The padded slots belong to + # ranks whose share is short by construction, and the unpacking loop + # below never reads them back. + stacked = tasks[0].new_empty((max_tasks, *tasks[0].shape)) + torch.stack(tasks, dim=0, out=stacked[: len(tasks)]) + stacked[len(tasks) :].zero_() + + key = (tuple(stacked.shape), stacked.dtype, str(stacked.device)) + buffer = self._tile_gather_buffers.get(key) + if buffer is None: + buffer = stacked.new_empty((sp_size, *stacked.shape)) + self._tile_gather_buffers[key] = buffer + self._tile_gather_stats["allocs"] += 1 + self._tile_gather_stats["bytes"] += buffer.numel() * buffer.element_size() + reuse = "alloc" + else: + self._tile_gather_stats["hits"] += 1 + reuse = "hit" + # Quantities, not just presence: a line that only says "installed" + # cannot distinguish a buffer that is being reused from one that is + # reallocated every request, which is the whole failure being fixed. + # Debug level, because this fires once per decoder tile batch (12 times + # per request on the canonical 1344x768 geometry) and the one-shot + # "installed" line above is what an operator needs at info level. + logger.debug( + "[H3_VAE_GATHER] reuse=%s key=%s/%s buf_ptr=0x%x hits=%d allocs=%d resident_mib=%.2f", + reuse, + tuple(stacked.shape), + stacked.dtype, + buffer.data_ptr(), + self._tile_gather_stats["hits"], + self._tile_gather_stats["allocs"], + self._tile_gather_stats["bytes"] / (1024.0 * 1024.0), + ) + dist.all_gather_into_tensor(buffer, stacked, group=group) + + results: list[torch.Tensor] = [None] * num_tiles # type: ignore[list-item] + for rank in range(sp_size): + num_rank_tasks = -(-(num_tiles - rank) // sp_size) + for k in range(num_rank_tasks): + results[k * sp_size + rank] = buffer[rank][k].clone() + return results + def _decoder_tile_count(self, latent: torch.Tensor) -> int: """Number of decoder tiles the checkpoint will split ``latent`` into.