anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Perf] Support TeaCache for MammothModa2 (#5357)

Browse Source 
Signed-off-by: nagisa-kun <1434936049@qq.com>
Signed-off-by: nagisa <1434936049@qq.com>
This commit is contained in:
hurukawa authored and GitHub committed 2026-10-05 01:25:56 +08:00
1 parent c9bee166fb
commit c9c478dd07
12 files changed
+684 -15
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/model_executor/models/mammoth_moda2/test_mammoth_moda2_teacache.py
+274
View File
@@ -0,0 +1,274 @@
from types import SimpleNamespace
from unittest.mock import Mock
import pytest
import torch
from torch import nn
from vllm_omni.diffusion.cache.cachedit import RequestScopedCacheDiTRuntime
from vllm_omni.diffusion.cache.teacache.config import TeaCacheConfig
from vllm_omni.diffusion.cache.teacache.extractors import extract_mammoth_moda2_context
from vllm_omni.diffusion.data import DiffusionOutput
from vllm_omni.diffusion.models.mammoth_moda2 import pipeline_mammothmoda2_dit as mammoth_pipeline_module
from vllm_omni.diffusion.models.mammoth_moda2.pipeline_mammothmoda2_dit import MammothModa2DiTPipeline
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
pytestmark = [pytest.mark.cpu]
def test_mammoth_moda2_calibrated_teacache_defaults():
config = TeaCacheConfig(transformer_type="MammothModa2Transformer2DModel")
assert config.rel_l1_thresh == 0.075
assert config.coefficients == [
-1761.242764481119,
859.398730831851,
-126.78044436937441,
9.16173963457422,
-0.1102826051200547,
]
class _FakeScheduler:
def set_timesteps(self, num_inference_steps, device, num_tokens): # noqa: ARG002
self.timesteps = torch.arange(num_inference_steps, 0, -1, device=device, dtype=torch.float32)
def step(self, model_pred, t, latents, return_dict=False): # noqa: ARG002
return (latents,)
class _FakeTransformer(nn.Module):
def __init__(self):
super().__init__()
self.config = SimpleNamespace(in_channels=4)
self.time_caption_embed = SimpleNamespace(image_embedder=None)
self.param = nn.Parameter(torch.zeros(()))
self.branches: list[str | None] = []
self.calls: list[dict] = []
def forward(self, hidden_states, **kwargs):
self.branches.append(kwargs.get("teacache_branch"))
self.calls.append(kwargs)
return torch.zeros_like(hidden_states)
class _FakeVAE(nn.Module):
def __init__(self):
super().__init__()
self.param = nn.Parameter(torch.zeros(()))
self.config = SimpleNamespace(scaling_factor=None, shift_factor=None)
def decode(self, latents, return_dict=False): # noqa: ARG002
return (latents,)
def _build_pipeline(monkeypatch):
monkeypatch.setattr(mammoth_pipeline_module, "FlowMatchEulerDiscreteScheduler", _FakeScheduler)
monkeypatch.setattr(
mammoth_pipeline_module,
"randn_tensor",
lambda shape, *, device, dtype, generator=None: torch.zeros(shape, device=device, dtype=dtype),
)
pipe = MammothModa2DiTPipeline.__new__(MammothModa2DiTPipeline)
nn.Module.__init__(pipe)
pipe.gen_transformer = _FakeTransformer()
pipe.gen_vae = _FakeVAE()
pipe.gen_image_condition_refiner = None
pipe.gen_freqs_cis = []
pipe._llm_hidden_size = 8
pipe._cache_dit_config = None
pipe._cache_dit_runtime = RequestScopedCacheDiTRuntime(pipe)
pipe.config = SimpleNamespace(
llm_config=SimpleNamespace(gen_vocab_start_index=100),
image_token_id=900,
video_token_id=901,
vision_start_token_id=902,
vision_end_token_id=903,
)
return pipe
def _run_pipeline(pipe, *, text_guidance_scale, cfg_range, num_inference_steps=4, num_requests=1):
request = OmniDiffusionRequest(
request_id="req-teacache",
prompt={
"prompt": "",
"height": 32,
"width": 32,
"additional_information": {
"full_hidden_states": torch.randn(3, 8),
"full_token_ids": [10, 11, 100],
"answer_start_index": 2,
},
},
sampling_params=OmniDiffusionSamplingParams(
height=32,
width=32,
guidance_scale=text_guidance_scale,
num_inference_steps=num_inference_steps,
extra_args={"cfg_range": cfg_range},
),
)
requests = [
OmniDiffusionRequest(
request_id=f"req-teacache-{i}",
prompt=request.prompt,
sampling_params=request.sampling_params,
)
for i in range(num_requests)
]
outputs = pipe(DiffusionRequestBatch(requests))
assert isinstance(outputs, list)
assert len(outputs) == num_requests
assert all(isinstance(output, DiffusionOutput) and output.output.shape == (1, 4, 4, 4) for output in outputs)
assert all(call["timestep"].shape == (num_requests,) for call in pipe.gen_transformer.calls)
return pipe.gen_transformer.branches
@pytest.mark.parametrize("num_requests", [1, 2])
def test_mammoth_moda2_non_cfg_passes_positive_teacache_branch(monkeypatch, num_requests):
pipe = _build_pipeline(monkeypatch)
branches = _run_pipeline(pipe, text_guidance_scale=1.0, cfg_range=[0.0, 1.0], num_requests=num_requests)
assert branches == ["positive", "positive", "positive", "positive"]
def test_mammoth_moda2_dev_forwards_ar_image_conditioning_to_positive_branch(monkeypatch):
"""Dev keeps AR image tokens separate and forwards them to the positive DiT call."""
pipe = _build_pipeline(monkeypatch)
pipe.gen_transformer.time_caption_embed.image_embedder = object()
_run_pipeline(
pipe,
text_guidance_scale=1.0,
cfg_range=[0.0, 1.0],
num_inference_steps=1,
)
assert len(pipe.gen_transformer.calls) == 1
call = pipe.gen_transformer.calls[0]
assert call["teacache_branch"] == "positive"
assert call["text_hidden_states"].shape == (1, 2, 8)
assert call["ar_image_hidden_states"].shape == (1, 1, 8)
assert torch.equal(call["ar_image_attention_mask"], torch.ones(1, 1, dtype=torch.bool))
@pytest.mark.parametrize("num_requests", [1, 2])
def test_mammoth_moda2_cfg_passes_positive_then_negative_teacache_branch(monkeypatch, num_requests):
pipe = _build_pipeline(monkeypatch)
branches = _run_pipeline(pipe, text_guidance_scale=4.0, cfg_range=[0.0, 1.0], num_requests=num_requests)
assert branches == [
"positive",
"negative",
"positive",
"negative",
"positive",
"negative",
"positive",
"negative",
]
@pytest.mark.parametrize("num_requests", [1, 2])
def test_mammoth_moda2_cfg_range_only_uses_negative_inside_range(monkeypatch, num_requests):
pipe = _build_pipeline(monkeypatch)
branches = _run_pipeline(pipe, text_guidance_scale=4.0, cfg_range=[0.5, 1.0], num_requests=num_requests)
assert branches == [
"positive",
"positive",
"positive",
"negative",
"positive",
"negative",
]
def test_mammoth_moda2_teacache_forwards_ar_image_conditioning():
"""The TeaCache extractor follows the full embedding contract, including the expanded text mask."""
hidden_size = 4
hidden_states = torch.zeros(1, hidden_size, 1, 1)
timestep = torch.zeros(1)
text_hidden_states = torch.zeros(1, 1, hidden_size)
text_attention_mask = torch.ones(1, 1, dtype=torch.bool)
freqs_cis = object()
ar_image_hidden_states = torch.randn(1, 2, hidden_size)
ar_image_attention_mask = torch.ones(1, 2, dtype=torch.bool)
temb = torch.zeros(1, hidden_size)
prepared_text_hidden_states = torch.randn(1, 3, hidden_size)
prepared_text_attention_mask = torch.ones(1, 3, dtype=torch.bool)
img_tokens = torch.randn(1, 1, hidden_size)
img_mask = torch.ones(1, 1, dtype=torch.bool)
context_rotary_emb = object()
noise_rotary_emb = object()
rotary_emb = object()
first_layer = SimpleNamespace(
norm1=Mock(side_effect=lambda joint_hidden_states, _temb: (joint_hidden_states,)),
)
module = SimpleNamespace(
layers=[first_layer],
config=SimpleNamespace(hidden_size=hidden_size, patch_size=1),
_validate_inputs=Mock(return_value=(1, 1, 1)),
_prepare_embeddings=Mock(
return_value=(
temb,
prepared_text_hidden_states,
prepared_text_attention_mask,
img_tokens,
img_mask,
1,
context_rotary_emb,
noise_rotary_emb,
rotary_emb,
[3],
[4],
)
),
_apply_refiners=Mock(return_value=(prepared_text_hidden_states, img_tokens)),
)
ctx = extract_mammoth_moda2_context(
module,
hidden_states=hidden_states,
timestep=timestep,
text_hidden_states=text_hidden_states,
text_attention_mask=text_attention_mask,
ref_image_hidden_states=None,
ar_image_hidden_states=ar_image_hidden_states,
ar_image_attention_mask=ar_image_attention_mask,
freqs_cis=freqs_cis,
teacache_branch="positive",
)
module._prepare_embeddings.assert_called_once_with(
hidden_states,
timestep,
text_hidden_states,
text_attention_mask,
freqs_cis,
1,
1,
1,
ar_image_hidden_states,
ar_image_attention_mask,
)
module._apply_refiners.assert_called_once_with(
prepared_text_hidden_states,
prepared_text_attention_mask,
context_rotary_emb,
img_tokens,
img_mask,
noise_rotary_emb,
temb,
)
assert ctx.hidden_states.shape == (1, 4, hidden_size)
assert ctx.extra_states["teacache_branch"] == "positive"