diff --git a/benchmarks/diffusion/benchmark_mammoth_moda2_qk_norm_rope.py b/benchmarks/diffusion/benchmark_mammoth_moda2_qk_norm_rope.py new file mode 100644 index 000000000..6117a7e12 --- /dev/null +++ b/benchmarks/diffusion/benchmark_mammoth_moda2_qk_norm_rope.py @@ -0,0 +1,262 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Benchmark MammothModa2's Q/K RMSNorm + adjacent-pair real RoPE. + +Example: + + CUDA_VISIBLE_DEVICES=0 python \ + benchmarks/diffusion/benchmark_mammoth_moda2_qk_norm_rope.py \ + --tokens 512 4173 8346 --warmup 50 --iters 200 +""" + +from __future__ import annotations + +import argparse +import importlib.metadata +import json +import os +import statistics +from collections.abc import Callable + +import torch + +from vllm_omni.diffusion.layers.fused_qk_norm_rope import ( + _fused_cuda_supported, + fused_qk_norm_rope, +) +from vllm_omni.diffusion.layers.norm import RMSNorm +from vllm_omni.diffusion.models.mammoth_moda2.rope_real import ( + apply_real_rotary_emb, +) + +_HEAD_DIM = 120 +_Q_HEADS = 21 +_KV_HEADS = 7 +_EPS = 1e-5 + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--tokens", nargs="+", type=int, default=[512, 4173, 8346]) + parser.add_argument("--warmup", type=int, default=50) + parser.add_argument("--iters", type=int, default=200) + parser.add_argument("--rounds", type=int, default=3) + parser.add_argument("--attention-sequences", nargs="*", type=int, default=[]) + return parser.parse_args() + + +def _version(distribution: str) -> str | None: + try: + return importlib.metadata.version(distribution) + except importlib.metadata.PackageNotFoundError: + return None + + +def _measure( + fn: Callable[[], object], + warmup: int, + iters: int, +) -> dict[str, float]: + for _ in range(warmup): + fn() + torch.accelerator.synchronize() + + torch.accelerator.reset_peak_memory_stats() + baseline_bytes = torch.accelerator.memory_allocated() + samples = [] + for _ in range(iters): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + fn() + end.record() + end.synchronize() + samples.append(float(start.elapsed_time(end))) + + ordered = sorted(samples) + return { + "median_ms": statistics.median(ordered), + "mean_ms": statistics.mean(ordered), + "stddev_ms": statistics.pstdev(ordered), + "min_ms": ordered[0], + "max_ms": ordered[-1], + "p90_ms": ordered[min(len(ordered) - 1, int(len(ordered) * 0.9))], + "extra_peak_mib": (torch.accelerator.max_memory_allocated() - baseline_bytes) / 2**20, + } + + +@torch.inference_mode() +def _run_shape(tokens: int, warmup: int, iters: int, *, fused_first: bool = False) -> dict[str, object]: + dtype = torch.bfloat16 + device = torch.device("cuda") + torch.manual_seed(42) + query = torch.randn(tokens, _Q_HEADS, _HEAD_DIM, device=device, dtype=dtype) + key = torch.randn(tokens, _KV_HEADS, _HEAD_DIM, device=device, dtype=dtype) + norm_q = RMSNorm(_HEAD_DIM, eps=_EPS).to(device=device, dtype=dtype) + norm_k = RMSNorm(_HEAD_DIM, eps=_EPS).to(device=device, dtype=dtype) + norm_q.weight.normal_(mean=1.0, std=0.2) + norm_k.weight.normal_(mean=1.0, std=0.2) + angles = torch.rand(tokens, _HEAD_DIM // 2, device=device) * (2 * torch.pi) + cos = angles.cos().repeat_interleave(2, dim=-1).to(dtype) + sin = angles.sin().repeat_interleave(2, dim=-1).to(dtype) + + def native() -> tuple[torch.Tensor, torch.Tensor]: + q_out = apply_real_rotary_emb(norm_q(query).unsqueeze(0), cos.unsqueeze(0), sin.unsqueeze(0)).squeeze(0) + k_out = apply_real_rotary_emb(norm_k(key).unsqueeze(0), cos.unsqueeze(0), sin.unsqueeze(0)).squeeze(0) + return q_out, k_out + + def fused() -> tuple[torch.Tensor, torch.Tensor]: + rope_table = torch.cat((cos[..., 0::2], sin[..., 0::2]), dim=-1) + return fused_qk_norm_rope( + query, + key, + norm_q.weight, + norm_k.weight, + rope_table, + _EPS, + head_dim=_HEAD_DIM, + rotary_dim=_HEAD_DIM, + interleaved=True, + ) + + expected_q, expected_k = native() + actual_q, actual_k = fused() + errors = { + "q_max_abs": (actual_q.float() - expected_q.float()).abs().max().item(), + "q_mean_abs": (actual_q.float() - expected_q.float()).abs().mean().item(), + "k_max_abs": (actual_k.float() - expected_k.float()).abs().max().item(), + "k_mean_abs": (actual_k.float() - expected_k.float()).abs().mean().item(), + } + torch.testing.assert_close(actual_q, expected_q, atol=0.0625, rtol=0.02) + torch.testing.assert_close(actual_k, expected_k, atol=0.0625, rtol=0.02) + + functions = {"native": native, "fused": fused} + order = ("fused", "native") if fused_first else ("native", "fused") + measured = {name: _measure(functions[name], warmup, iters) for name in order} + native_stats, fused_stats = measured["native"], measured["fused"] + return { + "tokens": tokens, + "native": native_stats, + "fused_including_rope_pack": fused_stats, + "median_speedup": native_stats["median_ms"] / fused_stats["median_ms"], + "errors": errors, + } + + +@torch.inference_mode() +def _run_attention(seq: int, warmup: int, iters: int, *, fused_first: bool = False) -> dict[str, object]: + from vllm_omni.diffusion.models.mammoth_moda2.mammothmoda2_dit_model import TransformerBlock + + torch.manual_seed(42) + block = ( + TransformerBlock( + 2520, + _Q_HEADS, + _KV_HEADS, + multiple_of=256, + ffn_dim_multiplier=1.0, + norm_eps=_EPS, + rope_repeats_pairs=True, + ) + .to(device="cuda", dtype=torch.bfloat16) + .eval() + ) + block.attn.norm_q.weight.normal_(1.0, 0.2) + block.attn.norm_k.weight.normal_(1.0, 0.2) + hidden = torch.randn(2, seq, 2520, device="cuda", dtype=torch.bfloat16) + mask = torch.ones(2, seq, device="cuda", dtype=torch.bool) + mask[0, seq - min(300, seq // 2) :] = False + mask[1, seq - min(17, seq // 2) :] = False + angles = torch.rand(2, seq, _HEAD_DIM // 2, device="cuda") * (2 * torch.pi) + # The production producer supplies FP32 tables; each example has its own positions. + rotary = angles.cos().repeat_interleave(2, dim=-1), angles.sin().repeat_interleave(2, dim=-1) + env = "VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS" + previous = os.environ.get(env) + + def run(): + return block.attn(hidden, hidden, attention_mask=mask, image_rotary_emb=rotary) + + try: + outputs = {} + timings = {} + order = ("fused", "native") if fused_first else ("native", "fused") + for name in order: + os.environ[env] = "0" if name == "fused" else str(2 * seq + 1) + outputs[name] = run() + timings[name] = _measure(run, warmup, iters) + diff = (outputs["native"].float() - outputs["fused"].float()).abs()[mask] + max_diff, mean_diff = diff.max().item(), diff.mean().item() + assert max_diff < 2e-2 and mean_diff < 1e-3, (max_diff, mean_diff) + assert torch.count_nonzero(outputs["fused"][~mask]) == 0 + return { + "sequence": seq, + "batch": 2, + "backend": block.attn.omni_attn.attn_backend.__name__, + "native": timings["native"], + "fused": timings["fused"], + "median_speedup": timings["native"]["median_ms"] / timings["fused"]["median_ms"], + "max_abs_diff": max_diff, + "mean_abs_diff": mean_diff, + } + finally: + if previous is None: + os.environ.pop(env, None) + else: + os.environ[env] = previous + + +def main() -> None: + args = _parse_args() + if ( + args.rounds < 1 + or args.iters < 1 + or args.warmup < 0 + or any(value < 1 for value in [*args.tokens, *args.attention_sequences]) + ): + raise ValueError("rounds, iters and sequence lengths must be positive; warmup must be nonnegative") + if not torch.cuda.is_available(): + raise RuntimeError("This benchmark requires CUDA") + + probe_q = torch.empty(1, _Q_HEADS, _HEAD_DIM, device="cuda", dtype=torch.bfloat16) + probe_k = torch.empty(1, _KV_HEADS, _HEAD_DIM, device="cuda", dtype=torch.bfloat16) + if not _fused_cuda_supported(probe_q, probe_k, _HEAD_DIM, _HEAD_DIM, interleaved=True): + raise RuntimeError("The fused CUDA QK norm/RoPE path is unavailable") + + result = { + "environment": { + "gpu": torch.cuda.get_device_name(), + "torch": torch.__version__, + "cuda": torch.version.cuda, + "triton": _version("triton"), + "vllm": _version("vllm"), + "vllm_omni": _version("vllm-omni"), + }, + "config": { + "dtype": "bfloat16", + "q_heads": _Q_HEADS, + "kv_heads": _KV_HEADS, + "head_dim": _HEAD_DIM, + "warmup": args.warmup, + "iters": args.iters, + "rounds": args.rounds, + "inference_mode": True, + "alternating_order": True, + "norm_weights": "normal(mean=1.0, std=0.2)", + }, + "results": [ + {"round": round_id, **_run_shape(tokens, args.warmup, args.iters, fused_first=bool(round_id % 2))} + for round_id in range(args.rounds) + for tokens in args.tokens + ], + "attention_results": [ + {"round": round_id, **_run_attention(seq, args.warmup, args.iters, fused_first=bool(round_id % 2))} + for round_id in range(args.rounds) + for seq in args.attention_sequences + ], + } + print(json.dumps(result, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/docs/configuration/environment_variables.md b/docs/configuration/environment_variables.md index a872eaf3c..9cc0abcad 100644 --- a/docs/configuration/environment_variables.md +++ b/docs/configuration/environment_variables.md @@ -130,7 +130,7 @@ settings. | Name | Type and default | Applies to and read time | Precedence and invalid values | Lifecycle | | --- | --- | --- | --- | --- | -| `VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS` | Non-negative integer; unset or empty means each consumer's own default (Boogu-Image: `2048`) | Token gate of the shared `fused_qk_norm_rope` op: a rotary table is packed for the fused kernel only when it spans at least this many positions (`B*S`); read at forward time | Environment-only. `0` = always fuse; a very large value disables the fused path. Any other value raises `ValueError` on the first forward. The crossover is host-dependent; re-benchmark before overriding. | Experimental performance control | +| `VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS` | Non-negative integer; unset or empty means each consumer's own default (Boogu-Image: `2048`; MammothModa2: `0`) | Token gate of the shared `fused_qk_norm_rope` op: a rotary table is packed for the fused kernel only when it spans at least this many positions (`B*S`); read at forward time | Environment-only. `0` = always fuse; a very large value disables the fused path. Any other value raises `ValueError` on the first forward. The crossover is host-dependent; re-benchmark before overriding. | Experimental performance control | | `VLLM_OMNI_SKIP_NVFP4_NAN_CLAMP` | Boolean truthy spellings: `1`, `true`, `yes`, `on`; default false | ModelOpt NVFP4 compatibility patch; read when `vllm_omni.patch` imports | Environment-only escape hatch. Any other value means false. Set only to diagnose the upstream NaN-scale issue. | Diagnostic and temporary | | `VLLM_OMNI_USE_QUACK_FP8` | Boolean truthy spellings: `1`, `true`, `yes`, `on`; unset means hardware auto-detection | FP8 scaled matrix multiplication; evaluated when quack capability is selected | A set value overrides auto-detection. Any non-truthy value forces quack off. If quack cannot load, vLLM-Omni warns and falls back to FlashInfer. | Experimental performance control | diff --git a/tests/diffusion/models/mammoth_moda2/test_dit_attention.py b/tests/diffusion/models/mammoth_moda2/test_dit_attention.py index bab2fee7d..49623d05c 100644 --- a/tests/diffusion/models/mammoth_moda2/test_dit_attention.py +++ b/tests/diffusion/models/mammoth_moda2/test_dit_attention.py @@ -187,3 +187,69 @@ def test_empty_text_stream_skips_the_kernel(monkeypatch, with_mask): image_rotary_emb=(angles.cos(), angles.sin()), ) assert out.shape == (BATCH, 0, DIM) + + +def test_model_owns_repeated_pair_rope_contract(): + from vllm_omni.diffusion.models.mammoth_moda2.mammothmoda2_dit_model import Transformer2DModel + from vllm_omni.diffusion.models.mammoth_moda2.rope_real import RotaryPosEmbedReal + + with set_current_diffusion_config(_SDPA_CONFIG): + model = Transformer2DModel( + hidden_size=48, + num_attention_heads=6, + num_kv_heads=2, + num_layers=1, + num_refiner_layers=1, + multiple_of=8, + ffn_dim_multiplier=1.0, + axes_dim_rope=(2, 2, 4), + axes_lens=(8, 8, 8), + text_feat_dim=32, + ) + for blocks in (model.layers, model.noise_refiner, model.ref_image_refiner, model.context_refiner): + assert all(block.attn.processor.rope_repeats_pairs for block in blocks) + assert not _block(2).attn.processor.rope_repeats_pairs + freqs = RotaryPosEmbedReal.get_freqs_real((2, 2, 4), (8, 8, 8), 10000) + ids = torch.tensor([[[0, 1, 2], [3, 4, 5]], [[1, 0, 3], [2, 5, 4]]]) + cos, sin = model.rope_embedder._get_freqs_real(freqs, ids) + for table in (cos, sin): + torch.testing.assert_close(table[..., 0::2], table[..., 1::2], atol=0, rtol=0) + assert not torch.equal(cos[0], cos[1]) + + +def test_embedding_preparation_preserves_stubbed_rope_tables(monkeypatch): + from types import SimpleNamespace + + from torch import nn + + import vllm_omni.diffusion.models.mammoth_moda2.mammothmoda2_dit_model as mammoth_dit + + # Embedding-only tests can stub the RoPE producer without constructing + # attention layers or supplying full transformer configuration. + model = object.__new__(mammoth_dit.Transformer2DModel) + nn.Module.__init__(model) + object.__setattr__(model, "_internal_dict", SimpleNamespace(patch_size=2)) + + class CaptionEmbedding(nn.Module): + def forward(self, timestep, text, dtype): + return torch.zeros(2, 8), text + + model.time_caption_embed = CaptionEmbedding() + model.x_embedder = nn.Identity() + model.rope_embedder = lambda *args: (None, None, None, None, [4, 4], [8, 8]) + + def forbid_gate(*args): + raise AssertionError("Embedding preparation must not resolve the fusion token gate") + + monkeypatch.setattr(mammoth_dit, "fused_qk_norm_rope_min_tokens", forbid_gate) + prepared = model._prepare_embeddings( + torch.zeros(2, 4, 4, 4), + torch.ones(2), + torch.zeros(2, 4, 8), + torch.ones(2, 4, dtype=torch.bool), + torch.zeros(1), + 2, + 4, + 4, + ) + assert prepared[6:9] == (None, None, None) diff --git a/tests/diffusion/models/mammoth_moda2/test_dit_attention_cuda.py b/tests/diffusion/models/mammoth_moda2/test_dit_attention_cuda.py index 050d404bd..cec4fa316 100644 --- a/tests/diffusion/models/mammoth_moda2/test_dit_attention_cuda.py +++ b/tests/diffusion/models/mammoth_moda2/test_dit_attention_cuda.py @@ -1,14 +1,16 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""Real-shape check of the MammothModa2 DiT attention on CUDA: the shared -backend (FLASH_ATTN by default, head_dim 120) against the pre-change arithmetic -in bf16, with the padding mask the joint text+image sequence carries.""" +"""Real-shape checks for MammothModa2's fused QK norm/RoPE and shared CUDA +attention backend against the native shared RMSNorm path.""" import pytest import torch +from diffusers.models.attention_processor import Attention import vllm_omni.diffusion.attention.backends.sdpa as sdpa_backend +import vllm_omni.diffusion.models.mammoth_moda2.mammothmoda2_dit_model as mammoth_dit from tests.helpers.mark import hardware_marks +from vllm_omni.diffusion.layers.norm import RMSNorm from vllm_omni.diffusion.models.mammoth_moda2.mammothmoda2_dit_model import TransformerBlock from .test_dit_attention import _reference_attention @@ -24,11 +26,127 @@ pytestmark = [ DIM, HEADS, KV_HEADS = 2520, 21, 7 +def _norm_attention(head_dim, dtype): + attn = Attention(query_dim=head_dim * HEADS, heads=HEADS, kv_heads=KV_HEADS, dim_head=head_dim) + attn.norm_q = RMSNorm(head_dim) + attn.norm_k = RMSNorm(head_dim) + return attn.to(device="cuda", dtype=dtype) + + +def test_unpaired_rope_preserves_native_arithmetic(monkeypatch): + """A generic block accepts independently valued even and odd RoPE lanes.""" + torch.manual_seed(42) + q = torch.randn(2, 17, HEADS, 120, device="cuda", dtype=torch.bfloat16) + k = torch.randn(2, 17, KV_HEADS, 120, device="cuda", dtype=torch.bfloat16) + attn = _norm_attention(120, torch.bfloat16) + angles = torch.rand(2, 17, 120, device="cuda") * 6.283 + rotary = angles.cos().to(q.dtype), angles.sin().to(q.dtype) + monkeypatch.setenv("VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS", "0") + with torch.no_grad(): + got = mammoth_dit._apply_qk_norm_rope(attn, q, k, rotary) + expected = ( + mammoth_dit.apply_real_rotary_emb(attn.norm_q(q), *rotary), + mammoth_dit.apply_real_rotary_emb(attn.norm_k(k), *rotary), + ) + for actual, want in zip(got, expected): + torch.testing.assert_close(actual, want, atol=0, rtol=0) + + +def test_qk_norm_rope_benchmark_measures_inference(monkeypatch): + from benchmarks.diffusion import benchmark_mammoth_moda2_qk_norm_rope as benchmark + + calls = [] + + def measure(fn, warmup, iters): + outputs = fn() + assert not torch.is_grad_enabled() + assert all(not output.requires_grad for output in outputs) + calls.append(fn.__name__) + return {"median_ms": 1.0} + + monkeypatch.setattr(benchmark, "_measure", measure) + benchmark._run_shape(17, warmup=0, iters=1) + assert calls == ["native", "fused"] + + +def _production_rotary(batch, seq, head_dim, dtype): + """Mammoth's real RoPE repeats one angle across each adjacent pair.""" + angles = torch.rand(batch, seq, head_dim // 2, device="cuda") * 6.283 + return ( + angles.cos().repeat_interleave(2, dim=-1).to(dtype), + angles.sin().repeat_interleave(2, dim=-1).to(dtype), + ) + + +@pytest.mark.parametrize("force_native", [False, True]) +def test_production_rope_uses_fused_qk_norm_rope(monkeypatch, force_native): + torch.manual_seed(0) + seq = 512 + block = ( + TransformerBlock( + DIM, + HEADS, + KV_HEADS, + multiple_of=256, + ffn_dim_multiplier=1.0, + norm_eps=1e-5, + rope_repeats_pairs=True, + ) + .cuda() + .to(torch.bfloat16) + .eval() + ) + hidden = torch.randn(2, seq, DIM, device="cuda", dtype=torch.bfloat16) + mask = torch.ones(2, seq, dtype=torch.bool, device="cuda") + mask[0, -37:] = False + rotary = _production_rotary(1, seq, block.head_dim, torch.bfloat16) + calls = [] + fused_qk_norm_rope = mammoth_dit.fused_qk_norm_rope + + def record_fused(q, k, q_weight, k_weight, rope_table, eps, **kwargs): + calls.append((q.shape, k.shape, rope_table.shape, kwargs)) + return fused_qk_norm_rope(q, k, q_weight, k_weight, rope_table, eps, **kwargs) + + monkeypatch.setenv("VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS", "0") + if force_native: + monkeypatch.setattr(mammoth_dit, "_fused_cuda_supported", lambda *args, **kwargs: False) + monkeypatch.setattr(mammoth_dit, "fused_qk_norm_rope", record_fused) + with torch.no_grad(): + got = block.attn(hidden, hidden, attention_mask=mask, image_rotary_emb=rotary) + want = _reference_attention(block.attn, hidden, mask, rotary) + + supported = mammoth_dit._fused_cuda_supported(hidden, hidden, block.head_dim, block.head_dim, interleaved=True) + expected_calls = ( + [ + ( + torch.Size((2 * seq, HEADS, block.head_dim)), + torch.Size((2 * seq, KV_HEADS, block.head_dim)), + torch.Size((2 * seq, block.head_dim)), + {"head_dim": block.head_dim, "rotary_dim": block.head_dim, "interleaved": True}, + ) + ] + if supported + else [] + ) + assert calls == expected_calls + diff = (got.float() - want.float()).abs()[mask] + assert diff.max().item() < 2e-2, diff.max().item() + assert diff.mean().item() < 1e-3, diff.mean().item() + + @pytest.mark.parametrize("seq", [77 + 4096, 512], ids=["t2i_1024", "short"]) def test_real_shape_matches_previous_arithmetic_bf16(seq): torch.manual_seed(0) block = ( - TransformerBlock(DIM, HEADS, KV_HEADS, multiple_of=256, ffn_dim_multiplier=1.0, norm_eps=1e-5) + TransformerBlock( + DIM, + HEADS, + KV_HEADS, + multiple_of=256, + ffn_dim_multiplier=1.0, + norm_eps=1e-5, + rope_repeats_pairs=True, + ) .cuda() .to(torch.bfloat16) .eval() @@ -37,8 +155,7 @@ def test_real_shape_matches_previous_arithmetic_bf16(seq): mask = torch.ones(2, seq, dtype=torch.bool, device="cuda") mask[0, seq - 300 :] = False mask[1, seq - 17 :] = False - angles = torch.rand(1, seq, block.head_dim, device="cuda") * 6.283 - rotary = (angles.cos().to(torch.bfloat16), angles.sin().to(torch.bfloat16)) + rotary = _production_rotary(1, seq, block.head_dim, torch.bfloat16) with torch.no_grad(): got = block.attn( hidden_states=hidden, encoder_hidden_states=hidden, attention_mask=mask, image_rotary_emb=rotary @@ -88,7 +205,15 @@ def test_fp32_falls_back_to_sdpa_and_matches_reference(monkeypatch, seq, force_g torch.manual_seed(0) dt = torch.float32 block = ( - TransformerBlock(DIM, HEADS, KV_HEADS, multiple_of=256, ffn_dim_multiplier=1.0, norm_eps=1e-5) + TransformerBlock( + DIM, + HEADS, + KV_HEADS, + multiple_of=256, + ffn_dim_multiplier=1.0, + norm_eps=1e-5, + rope_repeats_pairs=True, + ) .cuda() .to(dt) .eval() @@ -122,7 +247,11 @@ def test_fp32_falls_back_to_sdpa_and_matches_reference(monkeypatch, seq, force_g calls.append((query.shape, key.shape, value.shape, query.dtype, kwargs)) return sdpa(query, key, value, **kwargs) + def reject_fused(*args, **kwargs): + raise AssertionError("FP32 must keep Mammoth's native QK norm/RoPE path") + with torch.no_grad(), monkeypatch.context() as patch: + patch.setattr(mammoth_dit, "fused_qk_norm_rope", reject_fused) patch.setattr(sdpa_backend, "can_sdpa_use_fused_gqa", check_gqa) patch.setattr(torch.nn.functional, "scaled_dot_product_attention", record_sdpa) patch.setattr(strategy, "pre_attention", prepare) @@ -147,3 +276,206 @@ def test_fp32_falls_back_to_sdpa_and_matches_reference(monkeypatch, seq, force_g assert torch.count_nonzero(got[~mask]) == 0 diff = (got.float() - want.float()).abs()[mask] assert diff.max().item() < 1e-3, diff.max().item() + + +@pytest.mark.parametrize("layout", ["unbatched", "broadcast", "per_example", "strided"]) +@pytest.mark.parametrize("table_dtype", [torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("force_native", [False, True]) +def test_paired_rope_with_learned_norm_weights(monkeypatch, layout, table_dtype, force_native): + torch.manual_seed(42) + batch, seq, head_dim = 2, 19, 120 + attn = _norm_attention(head_dim, torch.bfloat16) + with torch.no_grad(): + attn.norm_q.weight.normal_(1.0, 0.2) + attn.norm_k.weight.normal_(1.0, 0.2) + # Mimic views into packed QKV projections, including their nontrivial stride. + packed = torch.randn(batch, seq, HEADS + 2 * KV_HEADS, head_dim, device="cuda", dtype=torch.bfloat16) + q, k, _ = packed.split((HEADS, KV_HEADS, KV_HEADS), dim=2) + rotary = _production_rotary(batch if layout in ("per_example", "strided") else 1, seq, head_dim, table_dtype) + if layout == "unbatched": + rotary = tuple(t.squeeze(0) for t in rotary) + elif layout == "strided": + rotary = tuple(t.transpose(0, 1).contiguous().transpose(0, 1) for t in rotary) + assert not rotary[0].is_contiguous() + calls = [] + fused = mammoth_dit.fused_qk_norm_rope + + def record(*args, **kwargs): + calls.append(args[4].shape) + return fused(*args, **kwargs) + + monkeypatch.setattr(mammoth_dit, "fused_qk_norm_rope", record) + if force_native: + monkeypatch.setattr(mammoth_dit, "_fused_cuda_supported", lambda *args, **kwargs: False) + monkeypatch.setenv("VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS", "0") + with torch.no_grad(): + got = mammoth_dit._apply_qk_norm_rope(attn, q, k, rotary, rope_repeats_pairs=True) + expected = ( + mammoth_dit.apply_real_rotary_emb(attn.norm_q(q), *rotary), + mammoth_dit.apply_real_rotary_emb(attn.norm_k(k), *rotary), + ) + supported = mammoth_dit._fused_cuda_supported(q, k, head_dim, head_dim, interleaved=True) + assert calls == ([torch.Size((batch * seq, head_dim))] if supported else []) + for actual, want in zip(got, expected): + torch.testing.assert_close(actual, want.to(actual.dtype), atol=0.0625, rtol=0.02) + + +@pytest.mark.parametrize( + "reason", ["support_query", "token_gate", "no_q_norm", "no_k_norm", "epsilon", "fp16", "head_dim", "no_rope"] +) +def test_qk_norm_rope_unsupported_inputs_keep_native_chain(monkeypatch, reason): + torch.manual_seed(42) + dtype = torch.float16 if reason == "fp16" else torch.bfloat16 + dim = 260 if reason == "head_dim" else 120 + attn = _norm_attention(dim, dtype) + q = torch.randn(2, 19, HEADS, dim, device="cuda", dtype=dtype) + k = torch.randn(2, 19, KV_HEADS, dim, device="cuda", dtype=dtype) + rotary = _production_rotary(2, 19, dim, dtype) if reason != "no_rope" else None + if reason == "support_query": + monkeypatch.setattr(mammoth_dit, "_fused_cuda_supported", lambda *args, **kwargs: False) + elif reason == "no_q_norm": + attn.norm_q = None + elif reason == "no_k_norm": + attn.norm_k = None + elif reason == "epsilon": + attn.norm_k.variance_epsilon = 1e-3 + monkeypatch.setenv("VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS", "39" if reason == "token_gate" else "0") + + def reject(*args, **kwargs): + raise AssertionError(f"Unsupported {reason} must keep the native chain") + + monkeypatch.setattr(mammoth_dit, "fused_qk_norm_rope", reject) + with torch.no_grad(): + got = mammoth_dit._apply_qk_norm_rope(attn, q, k, rotary, rope_repeats_pairs=True) + expected_q = attn.norm_q(q) if attn.norm_q is not None else q + expected_k = attn.norm_k(k) if attn.norm_k is not None else k + if rotary is not None: + expected_q = mammoth_dit.apply_real_rotary_emb(expected_q, *rotary) + expected_k = mammoth_dit.apply_real_rotary_emb(expected_k, *rotary) + torch.testing.assert_close(got[0], expected_q, atol=0, rtol=0) + torch.testing.assert_close(got[1], expected_k, atol=0, rtol=0) + + +def test_attention_benchmark_single_token_and_restores_gate(monkeypatch): + from benchmarks.diffusion import benchmark_mammoth_moda2_qk_norm_rope as benchmark + + env = "VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS" + monkeypatch.setenv(env, "123") + result = benchmark._run_attention(1, warmup=0, iters=1) + assert result["sequence"] == 1 + assert benchmark.os.environ[env] == "123" + + +@pytest.mark.parametrize("packed", [False, True], ids=["gated_native", "reused_packed"]) +def test_prepared_rope_table_skips_per_layer_gate_and_packing(monkeypatch, packed): + torch.manual_seed(7) + batch, seq, dim = 2, 19, 120 + attn = _norm_attention(dim, torch.bfloat16) + q = torch.randn(batch, seq, HEADS, dim, device="cuda", dtype=torch.bfloat16) + k = torch.randn(batch, seq, KV_HEADS, dim, device="cuda", dtype=torch.bfloat16) + cos, sin = _production_rotary(batch, seq, dim, torch.float32) + table = torch.cat((cos[..., 0::2], sin[..., 0::2]), dim=-1).reshape(batch * seq, dim) if packed else None + rotary = cos, sin, table + calls = [] + fused = mammoth_dit.fused_qk_norm_rope + + def record(*args, **kwargs): + calls.append(args[4]) + return fused(*args, **kwargs) + + def forbid_gate(*args, **kwargs): + raise AssertionError("Prepared RoPE must not resolve the token gate per attention layer") + + monkeypatch.setattr(mammoth_dit, "fused_qk_norm_rope", record) + monkeypatch.setattr(mammoth_dit, "fused_qk_norm_rope_min_tokens", forbid_gate) + with torch.no_grad(): + for _ in range(2): + got = mammoth_dit._apply_qk_norm_rope(attn, q, k, rotary, rope_repeats_pairs=True) + expected = ( + mammoth_dit.apply_real_rotary_emb(attn.norm_q(q), cos, sin), + mammoth_dit.apply_real_rotary_emb(attn.norm_k(k), cos, sin), + ) + for actual, want in zip(got, expected): + torch.testing.assert_close( + actual, want.to(actual.dtype), atol=0.0625 if packed else 0, rtol=0.02 if packed else 0 + ) + supported = packed and mammoth_dit._fused_cuda_supported(q, k, dim, dim, interleaved=True) + assert calls == ([table, table] if supported else []) + + +@pytest.mark.parametrize("force_native", [False, True]) +def test_model_prepares_three_rope_tables_with_one_threshold_lookup(monkeypatch, force_native): + from vllm_omni.diffusion.models.mammoth_moda2.mammothmoda2_dit_model import Transformer2DModel + from vllm_omni.diffusion.models.mammoth_moda2.rope_real import RotaryPosEmbedReal + + torch.manual_seed(7) + model = ( + Transformer2DModel( + patch_size=2, + in_channels=4, + hidden_size=48, + num_layers=1, + num_refiner_layers=1, + num_attention_heads=6, + num_kv_heads=2, + multiple_of=8, + ffn_dim_multiplier=1.0, + axes_dim_rope=(2, 2, 4), + axes_lens=(8, 8, 8), + text_feat_dim=32, + ) + .cuda() + .to(torch.bfloat16) + .eval() + ) + hidden = torch.randn(2, 4, 4, 4, device="cuda", dtype=torch.bfloat16) + text = torch.randn(2, 4, 32, device="cuda", dtype=torch.bfloat16) + mask = torch.ones(2, 4, device="cuda", dtype=torch.bool) + freqs = RotaryPosEmbedReal.get_freqs_real((2, 2, 4), (8, 8, 8), 10000) + calls = [] + + def threshold(default): + calls.append(default) + return 9 + + if force_native: + monkeypatch.setattr(mammoth_dit, "_fused_cuda_supported", lambda *args, **kwargs: False) + supported = mammoth_dit._fused_cuda_supported(hidden, hidden, 8, 8, interleaved=True) + tables: list[mammoth_dit.RotaryEmbedding] = [] + refiners = model._apply_refiners + layers = model._apply_transformer_layers + + def record_refiners(text, text_mask, context, img, img_mask, noise, temb): + tables.extend((context, noise)) + return refiners(text, text_mask, context, img, img_mask, noise, temb) + + def record_layers(hidden, mask, joint, temb): + tables.append(joint) + return layers(hidden, mask, joint, temb) + + monkeypatch.setattr(model, "_apply_refiners", record_refiners) + monkeypatch.setattr(model, "_apply_transformer_layers", record_layers) + monkeypatch.setattr(mammoth_dit, "fused_qk_norm_rope_min_tokens", threshold) + args = (hidden, torch.ones(2, device="cuda"), text, freqs, mask) + with torch.no_grad(): + fused_output = model(*args) + context, noise, joint = tables + assert calls == ([0] if supported else []) + assert all(len(rotary) == 3 for rotary in (context, noise, joint)) + assert context[2] is None and noise[2] is None + if supported: + assert joint[2].shape == (2 * 8, 8) + else: + assert joint[2] is None + for rotary in (context, noise, joint): + if rotary[2] is not None: + expected = torch.cat((rotary[0][..., 0::2], rotary[1][..., 0::2]), dim=-1).reshape(-1, 8).float() + torch.testing.assert_close(rotary[2], expected, atol=0, rtol=0) + + # Exercise the prepared tuple through every real refiner and joint block, + # then compare the same model forward with the unsupported-device fallback. + with torch.no_grad(): + monkeypatch.setattr(mammoth_dit, "_fused_cuda_supported", lambda *args, **kwargs: False) + native_output = model(*args) + assert fused_output.shape == hidden.shape + torch.testing.assert_close(fused_output, native_output, atol=0.0625, rtol=0.02) diff --git a/vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py b/vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py index 4b5edf717..c6bdb5719 100644 --- a/vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py +++ b/vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py @@ -12,10 +12,112 @@ from torch import nn from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata from vllm_omni.diffusion.attention.layer import Attention as OmniAttention +from vllm_omni.diffusion.layers.fused_qk_norm_rope import ( + _fused_cuda_supported, + fused_qk_norm_rope, + fused_qk_norm_rope_min_tokens, +) from vllm_omni.diffusion.layers.norm import RMSNorm from .rope_real import RotaryPosEmbedReal, apply_real_rotary_emb +# On the measured RTX 4090, fusion won even for the shortest tested tables; +# deployments can tune the crossover with VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS. +_MAMMOTH_FUSED_QK_NORM_ROPE_MIN_TOKENS = 0 + +RotaryEmbedding = tuple[torch.Tensor, torch.Tensor] | tuple[torch.Tensor, torch.Tensor, torch.Tensor | None] + + +def _with_packed_rope_table(rotary_emb: tuple[torch.Tensor, torch.Tensor], min_tokens: int | None) -> RotaryEmbedding: + """Pack a model-produced adjacent-pair RoPE table once for all layers. + + A third element of ``None`` records that the support or token gate chose + the eager path, so attention layers need not resolve that gate again. + """ + cos, sin = rotary_emb + if min_tokens is None or cos.shape[0] * cos.shape[1] < min_tokens: + return cos, sin, None + packed = torch.cat((cos[..., 0::2], sin[..., 0::2]), dim=-1) + return cos, sin, packed.reshape(-1, cos.shape[-1]).float() + + +def _apply_qk_norm_rope( + attn: Attention, + query: torch.Tensor, + key: torch.Tensor, + image_rotary_emb: RotaryEmbedding | None, + *, + rope_repeats_pairs: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + """Apply Mammoth's per-head Q/K norm and adjacent-pair real RoPE.""" + batch_size, sequence_length, _, head_dim = query.shape + cos, sin = (image_rotary_emb[0], image_rotary_emb[1]) if image_rotary_emb is not None else (None, None) + prepared = image_rotary_emb is not None and len(image_rotary_emb) == 3 + packed = image_rotary_emb[2] if prepared else None + use_fused = ( + rope_repeats_pairs + and cos is not None + and sin is not None + and attn.norm_q is not None + and attn.norm_k is not None + and attn.norm_q.variance_epsilon == attn.norm_k.variance_epsilon + and key.shape[0] == batch_size + and key.shape[1] == sequence_length + and key.shape[-1] == head_dim + and _fused_cuda_supported(query, key, head_dim, head_dim, interleaved=True) + and cos.shape[-2:] == (sequence_length, head_dim) + and sin.shape == cos.shape + and cos.ndim in (2, 3) + and (cos.ndim == 2 or cos.shape[0] in (1, batch_size)) + and cos.device == query.device + and cos.dtype in (query.dtype, torch.float32) + and sin.device == query.device + and sin.dtype == cos.dtype + and ( + packed is not None + if prepared + else batch_size * sequence_length >= fused_qk_norm_rope_min_tokens(_MAMMOTH_FUSED_QK_NORM_ROPE_MIN_TOKENS) + ) + ) + if use_fused: + assert cos is not None and sin is not None + if prepared: + assert packed is not None + rope_table = packed + else: + if cos.ndim == 2: + cos, sin = cos.unsqueeze(0), sin.unsqueeze(0) + if cos.shape[0] == 1 and batch_size > 1: + cos = cos.expand(batch_size, -1, -1) + sin = sin.expand(batch_size, -1, -1) + rope_table = torch.cat((cos[..., 0::2], sin[..., 0::2]), dim=-1).reshape( + batch_size * sequence_length, head_dim + ) + query, key = fused_qk_norm_rope( + query.reshape(batch_size * sequence_length, query.shape[2], head_dim), + key.reshape(batch_size * sequence_length, key.shape[2], head_dim), + attn.norm_q.weight, + attn.norm_k.weight, + rope_table, + attn.norm_q.variance_epsilon, + head_dim=head_dim, + rotary_dim=head_dim, + interleaved=True, + ) + return ( + query.reshape(batch_size, sequence_length, query.shape[1], head_dim), + key.reshape(batch_size, sequence_length, key.shape[1], head_dim), + ) + + if attn.norm_q is not None: + query = attn.norm_q(query) + if attn.norm_k is not None: + key = attn.norm_k(key) + if image_rotary_emb is not None: + query = apply_real_rotary_emb(query, image_rotary_emb[0], image_rotary_emb[1]) + key = apply_real_rotary_emb(key, image_rotary_emb[0], image_rotary_emb[1]) + return query, key + class LuminaRMSNormZero(nn.Module): """ @@ -275,13 +377,18 @@ class AttnProcessor: head count first. """ + def __init__(self, *, rope_repeats_pairs: bool = False) -> None: + # Only the model's RotaryPosEmbedReal producer guarantees equal adjacent + # lanes. Generic processors retain the native four-lane arithmetic. + self.rope_repeats_pairs = rope_repeats_pairs + def __call__( self, attn: Attention, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None, - image_rotary_emb: torch.Tensor | None = None, + image_rotary_emb: RotaryEmbedding | None = None, ) -> torch.Tensor: batch_size, sequence_length, _ = hidden_states.shape @@ -307,14 +414,7 @@ class AttnProcessor: key = key.view(batch_size, -1, kv_heads, head_dim) value = value.view(batch_size, -1, kv_heads, head_dim) - if attn.norm_q is not None: - query = attn.norm_q(query) - if attn.norm_k is not None: - key = attn.norm_k(key) - - if image_rotary_emb is not None: - query = apply_real_rotary_emb(query, image_rotary_emb[0], image_rotary_emb[1]) - key = apply_real_rotary_emb(key, image_rotary_emb[0], image_rotary_emb[1]) + query, key = _apply_qk_norm_rope(attn, query, key, image_rotary_emb, rope_repeats_pairs=self.rope_repeats_pairs) query, key = query.to(dtype), key.to(dtype) @@ -347,13 +447,15 @@ class TransformerBlock(nn.Module): ffn_dim_multiplier: float, norm_eps: float, modulation: bool = True, + *, + rope_repeats_pairs: bool = False, ) -> None: """Initialize the transformer block.""" super().__init__() self.head_dim = dim // num_attention_heads self.modulation = modulation - processor = AttnProcessor() + processor = AttnProcessor(rope_repeats_pairs=rope_repeats_pairs) # Initialize attention layer self.attn = Attention( @@ -403,7 +505,7 @@ class TransformerBlock(nn.Module): self, hidden_states: torch.Tensor, attention_mask: torch.Tensor, - image_rotary_emb: torch.Tensor, + image_rotary_emb: RotaryEmbedding, temb: torch.Tensor | None = None, ) -> torch.Tensor: if self.modulation: @@ -505,6 +607,7 @@ class Transformer2DModel(ModelMixin, ConfigMixin): ffn_dim_multiplier, norm_eps, modulation=True, + rope_repeats_pairs=True, ) for _ in range(num_refiner_layers) ] @@ -520,6 +623,7 @@ class Transformer2DModel(ModelMixin, ConfigMixin): ffn_dim_multiplier, norm_eps, modulation=True, + rope_repeats_pairs=True, ) for _ in range(num_refiner_layers) ] @@ -535,6 +639,7 @@ class Transformer2DModel(ModelMixin, ConfigMixin): ffn_dim_multiplier, norm_eps, modulation=False, + rope_repeats_pairs=True, ) for _ in range(num_refiner_layers) ] @@ -551,6 +656,7 @@ class Transformer2DModel(ModelMixin, ConfigMixin): ffn_dim_multiplier, norm_eps, modulation=True, + rope_repeats_pairs=True, ) for _ in range(num_layers) ] @@ -705,10 +811,10 @@ class Transformer2DModel(ModelMixin, ConfigMixin): self, text_hidden_states: torch.Tensor, text_attention_mask: torch.Tensor, - context_rotary_emb: torch.Tensor, + context_rotary_emb: RotaryEmbedding, img_tokens: torch.Tensor, img_mask: torch.Tensor, - noise_rotary_emb: torch.Tensor, + noise_rotary_emb: RotaryEmbedding, temb: torch.Tensor, ): for layer in self.context_refiner: @@ -777,6 +883,20 @@ class Transformer2DModel(ModelMixin, ConfigMixin): ar_image_attention_mask, ) + # The producer repeats each frequency across adjacent lanes. Resolve + # support and the token gate once, then share each of the three packed + # tables across its refiner or transformer layers. Keep preparation in + # forward so embedding-only callers retain the producer's return values. + head_dim = rotary_emb[0].shape[-1] + min_tokens = ( + fused_qk_norm_rope_min_tokens(_MAMMOTH_FUSED_QK_NORM_ROPE_MIN_TOKENS) + if _fused_cuda_supported(hidden_states, hidden_states, head_dim, head_dim, interleaved=True) + else None + ) + context_rotary_emb = _with_packed_rope_table(context_rotary_emb, min_tokens) + noise_rotary_emb = _with_packed_rope_table(noise_rotary_emb, min_tokens) + rotary_emb = _with_packed_rope_table(rotary_emb, min_tokens) + text_hidden_states, img_tokens = self._apply_refiners( text_hidden_states, text_attention_mask,