diff --git a/.gitignore b/.gitignore index 485b7d415..75b3c67f2 100644 --- a/.gitignore +++ b/.gitignore @@ -4,6 +4,9 @@ compile_commands.json # Visual Studio Code configs. .vscode/ +# Neovim configs. +.nvim.lua + .pytest-tmp* # Byte-compiled / optimized / DLL files @@ -93,6 +96,7 @@ celerybeat-schedule .env .venv .venv* +.direnv/ env/ venv/ ENV/ diff --git a/AGENTS.md b/AGENTS.md index 1c4624c1b..4b903fc86 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -1,3 +1,34 @@ +# Principle Agent Guide + +You are a lazy senior developer. Lazy means efficient, not careless. The best code is the code never written. + +Before writing any code, stop at the first rung that holds: + +1. Does this need to be built at all? (YAGNI) +2. Does it already exist in this codebase? Reuse the helper, util, or pattern that's already here, don't re-write it. +3. Does the standard library already do this? Use it. +4. Does a native platform feature cover it? Use it. +5. Does an already-installed dependency solve it? Use it. +6. Can this be one line? Make it one line. +7. Only then: write the minimum code that works. + +The ladder runs after you understand the problem, not instead of it: read the task and the code it touches, trace the real flow end to end, then climb. + +Bug fix = root cause, not symptom: a report names a symptom. Grep every caller of the function you touch and fix the shared function once — one guard there is a smaller diff than one per caller, and patching only the path the ticket names leaves a sibling caller still broken. + +Rules: + +- No abstractions that weren't explicitly requested. +- No new dependency if it can be avoided. +- No boilerplate nobody asked for. +- Deletion over addition. Boring over clever. Fewest files possible. +- Shortest working diff wins, but only once you understand the problem. The smallest change in the wrong place isn't lazy, it's a second bug. +- Question complex requests: "Do you actually need X, or does Y cover it?" +- Pick the edge-case-correct option when two stdlib approaches are the same size, lazy means less code, not the flimsier algorithm. +- Mark deliberate simplifications that cut a real corner with a known ceiling (global lock, O(n²) scan, naive heuristic) with a `ponytail:` comment naming the ceiling and upgrade path. + +Not lazy about: understanding the problem (read it fully and trace the real flow before picking a rung, a small diff you don't understand is just laziness dressed up as efficiency), input validation at trust boundaries, error handling that prevents data loss, security, accessibility, the calibration real hardware needs (the platform is never the spec ideal, a clock drifts, a sensor reads off), anything explicitly requested. Lazy code without its check is unfinished: non-trivial logic leaves ONE runnable check behind, the smallest thing that fails if the logic breaks (an assert-based demo/self-check or one small test file; no frameworks, no fixtures). Trivial one-liners need no test. + # FlashDreams Agent Guide FlashDreams is a GPU-heavy inference and serving library for autoregressive video and world models. Default to inspection, docs, config checks, and CPU tests unless the user explicitly asks to run generation or GPU workflows. diff --git a/THIRD-PARTY-NOTICES b/THIRD-PARTY-NOTICES index 433a0c0e2..7e4a84fc5 100644 --- a/THIRD-PARTY-NOTICES +++ b/THIRD-PARTY-NOTICES @@ -33,6 +33,8 @@ huggingface-hub Apache-2.0 https://github.com/huggingface/huggingface_h loguru MIT https://github.com/Delgan/loguru numpy BSD-3-Clause https://github.com/numpy/numpy nvidia-ml-py BSD-3-Clause https://pypi.org/project/nvidia-ml-py/ +nvtx Apache-2.0 WITH LLVM-exception + https://github.com/NVIDIA/NVTX psutil BSD-3-Clause https://github.com/giampaolo/psutil safetensors Apache-2.0 https://github.com/huggingface/safetensors torch BSD-3-Clause https://pytorch.org diff --git a/flashdreams/benchmarks/accelerated/test_multi_head_attention_benchmark.py b/flashdreams/benchmarks/accelerated/test_multi_head_attention_benchmark.py new file mode 100644 index 000000000..5e0b00fb3 --- /dev/null +++ b/flashdreams/benchmarks/accelerated/test_multi_head_attention_benchmark.py @@ -0,0 +1,503 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Microbenchmarks for Triton multi-head attention. + +All cases use identical geometry and deterministic random weights. + +Run the manual GPU benchmarks with:: + + uv run --package flashdreams --group test pytest \ + flashdreams/benchmarks/accelerated/test_multi_head_attention_benchmark.py \ + -p no:manual_marker -m manual --benchmark-only -v +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import pytest +import torch +from pytest_benchmark.fixture import BenchmarkFixture +from torch import Tensor + +from flashdreams.accelerated.multi_head_attention import AttentionType, QKNormScope +from flashdreams.accelerated.multi_head_attention_triton import ( + QKVFusionOption, + SDPABackend, + TritonMultiHeadAttention, +) + +pytestmark = [ + pytest.mark.manual, + pytest.mark.skipif( + not torch.cuda.is_available(), + reason="Accelerated multi-head attention benchmarks require CUDA.", + ), +] + +_WARMUP_ROUNDS = 5 +"""Warmup calls used to absorb kernel compilation and autotuning.""" + +_BENCHMARK_ROUNDS = 50 +"""Measured calls used for each benchmark case.""" + + +@dataclass(frozen=True) +class _ImplementationCase: + """One Triton configuration within a shared benchmark group.""" + + sdpa_backend: SDPABackend + """Triton attention backend.""" + + qkv_fusion_option: QKVFusionOption + """Triton projection fusion policy.""" + + use_fp8: bool = False + """Whether the case enables FP8 projections and supported storage.""" + + @property + def id(self) -> str: + """Return a stable pytest identifier for this implementation row.""" + backend = "fa2" if self.sdpa_backend is SDPABackend.TRITON else "cudnn" + precision = "fp8" if self.use_fp8 else "bf16" + fusion = self.qkv_fusion_option.value.replace("_", "-") + return f"triton-{backend}-{precision}-{fusion}" + + +_IMPLEMENTATION_CASES = tuple( + _ImplementationCase( + sdpa_backend=sdpa_backend, + qkv_fusion_option=qkv_fusion_option, + use_fp8=use_fp8, + ) + for sdpa_backend in SDPABackend + for use_fp8 in (False, True) + for qkv_fusion_option in QKVFusionOption +) +"""Every Triton backend, precision, and fusion row.""" + +_SHARED_CONFIGS = tuple( + pytest.param( + qk_norm_scope, + rope_interleaved, + bias, + id=( + f"norm-{qk_norm_scope.value}-" + f"rope-{'interleaved' if rope_interleaved else 'split'}-" + f"bias-{'on' if bias else 'off'}" + ), + ) + for qk_norm_scope in QKNormScope + for rope_interleaved in (False, True) + for bias in (False, True) +) +"""Policies shared by every Triton case, each mapped to its own benchmark group.""" + + +class _TritonMultiHeadAttention(TritonMultiHeadAttention): + """Canonical Triton attention implementation used by benchmarks.""" + + @property + def query_projection(self) -> torch.nn.Linear: + """Return the canonical query projection.""" + return self.q_proj + + @property + def key_projection(self) -> torch.nn.Linear: + """Return the canonical key projection.""" + return self.k_proj + + @property + def value_projection(self) -> torch.nn.Linear: + """Return the canonical value projection.""" + return self.v_proj + + @property + def output_projection(self) -> torch.nn.Linear: + """Return the canonical output projection.""" + return self.output_proj + + @property + def query_norm(self) -> torch.nn.Module: + """Return the canonical query normalization.""" + return self.q_norm + + @property + def key_norm(self) -> torch.nn.Module: + """Return the canonical key normalization.""" + return self.k_norm + + def __init__( + self, + query_dim: int, + n_heads: int = 8, + head_dim: int = 64, + *, + context_dim: int | None = None, + attention_type: AttentionType = AttentionType.SELF_ATTENTION, + qkv_fusion_option: QKVFusionOption = QKVFusionOption.FULL, + qkv_bias: bool = False, + output_bias: bool = False, + qk_norm_scope: QKNormScope = QKNormScope.HEAD, + rope_interleaved: bool = False, + use_fp8: bool = False, + sdpa_backend: SDPABackend = SDPABackend.CUDNN, + ) -> None: + """Initialize canonical projections and normalization modules.""" + super().__init__( + query_dim=query_dim, + n_heads=n_heads, + head_dim=head_dim, + context_dim=context_dim, + attention_type=attention_type, + qkv_fusion_option=qkv_fusion_option, + qk_norm_scope=qk_norm_scope, + rope_interleaved=rope_interleaved, + use_fp8=use_fp8, + sdpa_backend=sdpa_backend, + ) + self.q_proj = torch.nn.Linear(self.query_dim, self.inner_dim, bias=qkv_bias) + self.k_proj = torch.nn.Linear(self.context_dim, self.inner_dim, bias=qkv_bias) + self.v_proj = torch.nn.Linear(self.context_dim, self.inner_dim, bias=qkv_bias) + self.output_proj = torch.nn.Linear( + self.inner_dim, self.query_dim, bias=output_bias + ) + if self.qk_norm_scope is QKNormScope.NONE: + self.q_norm = torch.nn.Identity() + self.k_norm = torch.nn.Identity() + else: + norm_dim = ( + self.head_dim + if self.qk_norm_scope is QKNormScope.HEAD + else self.inner_dim + ) + self.q_norm = torch.nn.RMSNorm(norm_dim, eps=self.qk_norm_eps) + self.k_norm = torch.nn.RMSNorm(norm_dim, eps=self.qk_norm_eps) + self._initialize_derived_weights() + + +_BATCH_SIZE = 1 +_DTYPE = torch.bfloat16 +_SEED = 42 +_SINK_SIZE = 0 + +_QUERY_DIM = 2048 +"""Input and output feature width shared by all cases.""" + +_N_HEADS = 16 +"""Number of attention heads shared by all cases.""" + +_HEAD_DIM = _QUERY_DIM // _N_HEADS + +_CHUNK_SIZE = 80 * 60 +"""Number of query tokens processed by each benchmark call.""" + +_WINDOW_CHUNKS = 6 +"""Number of chunks retained in the full rolling cache.""" + +_WINDOW_SIZE = _WINDOW_CHUNKS * _CHUNK_SIZE + + +def _make_attention( + case: _ImplementationCase, + *, + attention_type: AttentionType, + qk_norm_scope: QKNormScope, + rope_interleaved: bool, + bias: bool, +) -> _TritonMultiHeadAttention: + """Build one deterministic Triton benchmark case.""" + return _TritonMultiHeadAttention( + query_dim=_QUERY_DIM, + context_dim=_QUERY_DIM, + n_heads=_N_HEADS, + head_dim=_HEAD_DIM, + attention_type=attention_type, + qkv_fusion_option=case.qkv_fusion_option, + qkv_bias=bias, + output_bias=bias, + qk_norm_scope=qk_norm_scope, + rope_interleaved=rope_interleaved, + use_fp8=case.use_fp8, + sdpa_backend=case.sdpa_backend, + ) + + +@torch.inference_mode() +def _benchmark_multi_head_attention( + benchmark: BenchmarkFixture, + case: _ImplementationCase, + *, + attention_type: AttentionType, + qk_norm_scope: QKNormScope, + rope_interleaved: bool, + bias: bool, +) -> None: + """Run one synchronized attention benchmark within a shared-policy group. + + Streaming self-attention times forward over a prefilled rolling cache. + Cross-attention times static K/V preparation and forward together so the + requested fusion variants exercise the work they actually change. + + Args: + benchmark: Pytest benchmark fixture used to record synchronized timings. + case: Implementation-specific backend, precision, and fusion settings. + attention_type: Self- or cross-attention benchmark family. + qk_norm_scope: Shared Q/K normalization policy. + rope_interleaved: Shared rotary-pair layout. + bias: Whether every Q/K/V and output projection uses a bias. + """ + if not torch.cuda.is_bf16_supported(): + pytest.skip("Multi-head attention benchmark requires bfloat16 support.") + + device = torch.device("cuda") + if torch.cuda.get_device_capability(device)[0] < 9: + pytest.skip( + "Triton accelerated attention requires compute capability 9.0 or newer." + ) + + torch.manual_seed(_SEED) + attention = _make_attention( + case, + attention_type=attention_type, + qk_norm_scope=qk_norm_scope, + rope_interleaved=rope_interleaved, + bias=bias, + ) + attention.to(device=device, dtype=_DTYPE).eval() + + generator = torch.Generator(device=device).manual_seed(_SEED) + inputs = [ + torch.randn( + _BATCH_SIZE, + _CHUNK_SIZE, + _QUERY_DIM, + generator=generator, + device=device, + dtype=_DTYPE, + ) + for _ in range(_WINDOW_CHUNKS + 1) + ] + rope_freqs = [ + torch.randn( + _CHUNK_SIZE, + 1, + 1, + _HEAD_DIM, + generator=generator, + device=device, + dtype=torch.float32, + ) + for _ in range(_WINDOW_CHUNKS + 1) + ] + + attention_label = ( + "self" if attention_type is AttentionType.SELF_ATTENTION else "cross" + ) + rope_label = "interleaved" if rope_interleaved else "split" + bias_label = "on" if bias else "off" + benchmark.group = "-".join( + ( + "multi-head-attention", + attention_label, + "norm", + qk_norm_scope.value, + "rope", + rope_label, + "bias", + bias_label, + ) + ) + + cache_dtype = ( + torch.float8_e4m3fn + if case.use_fp8 and case.sdpa_backend is SDPABackend.TRITON + else _DTYPE + ) + benchmark.extra_info.update( + { + "implementation": "triton", + "implementation_case": case.id, + "attention_type": attention_type.value, + "timed_region": ( + "forward" + if attention_type is AttentionType.SELF_ATTENTION + else "compute_kv_and_forward" + ), + "batch_size": _BATCH_SIZE, + "query_tokens": _CHUNK_SIZE, + "context_tokens": _WINDOW_SIZE, + "query_dim": _QUERY_DIM, + "context_dim": _QUERY_DIM, + "num_heads": _N_HEADS, + "head_dim": _HEAD_DIM, + "parameter_count": sum( + parameter.numel() for parameter in attention.parameters() + ), + "checkpoint": "random_init_shared_weights", + "dtype": str(_DTYPE), + "cache_dtype": str(cache_dtype), + "bias": bias, + "qkv_bias": bias, + "output_bias": bias, + "qk_norm_scope": qk_norm_scope.value, + "rope_interleaved": rope_interleaved, + "sdpa_backend": case.sdpa_backend.value, + "qkv_fusion_option": case.qkv_fusion_option.value, + "use_fp8": case.use_fp8, + "cache_state": ( + "full_rolling_window" + if attention_type is AttentionType.SELF_ATTENTION + else "rebuilt_static_context" + ), + "cache_prefill_chunks": ( + _WINDOW_CHUNKS if attention_type is AttentionType.SELF_ATTENTION else 0 + ), + "gpu": torch.cuda.get_device_name(device), + "torch": torch.__version__, + "cuda": torch.version.cuda, + "cudnn": torch.backends.cudnn.version(), + "warmup_rounds": _WARMUP_ROUNDS, + "benchmark_rounds": _BENCHMARK_ROUNDS, + "seed": _SEED, + } + ) + + query = inputs[_WINDOW_CHUNKS] + query_rope = rope_freqs[_WINDOW_CHUNKS] + if attention_type is AttentionType.SELF_ATTENTION: + cache = attention.allocate_kv_cache( + batch_size=_BATCH_SIZE, + chunk_size=_CHUNK_SIZE, + window_size=_WINDOW_SIZE, + sink_size=_SINK_SIZE, + device=device, + dtype=_DTYPE, + ) + + # Prefill and roll outside the timer. Every measured forward sees the + # same full context and overwrites the same final cache slot. + for chunk_idx in range(_WINDOW_CHUNKS): + cache.before_update(chunk_idx) + attention(inputs[chunk_idx], cache, rope_freqs[chunk_idx]) + cache.after_update(chunk_idx) + cache.before_update(_WINDOW_CHUNKS) + torch.cuda.synchronize() + + def synchronized_self_forward() -> Tensor: + result = attention(query, cache, query_rope) + torch.cuda.synchronize() + return result + + output = benchmark.pedantic( + synchronized_self_forward, + iterations=1, + rounds=_BENCHMARK_ROUNDS, + warmup_rounds=_WARMUP_ROUNDS, + ) + cache.after_update(_WINDOW_CHUNKS) + else: + context = torch.cat(inputs[:_WINDOW_CHUNKS], dim=1) + context_rope = torch.cat(rope_freqs[:_WINDOW_CHUNKS], dim=0) + torch.cuda.synchronize() + + # Static K/V projection is part of this end-to-end cross-attention + # measurement because fusion changes that stage, not query-only forward. + def synchronized_cross_forward() -> Tensor: + cache = attention.compute_kv(context, context_rope) + result = attention(query, cache, query_rope) + torch.cuda.synchronize() + return result + + output = benchmark.pedantic( + synchronized_cross_forward, + iterations=1, + rounds=_BENCHMARK_ROUNDS, + warmup_rounds=_WARMUP_ROUNDS, + ) + + assert output.shape == query.shape + assert torch.isfinite(output).all() + + +@pytest.mark.parametrize( + "case", + _IMPLEMENTATION_CASES, + ids=lambda case: case.id, +) +@pytest.mark.parametrize( + "qk_norm_scope,rope_interleaved,bias", + _SHARED_CONFIGS, +) +def test_self_attention_benchmark( + benchmark: BenchmarkFixture, + case: _ImplementationCase, + qk_norm_scope: QKNormScope, + rope_interleaved: bool, + bias: bool, +) -> None: + """Benchmark streaming self-attention within one shared-policy group. + + Args: + benchmark: Pytest benchmark fixture used to record synchronized timings. + case: Triton backend, precision, and fusion row. + qk_norm_scope: Shared Q/K normalization policy defining the group. + rope_interleaved: Shared rotary-pair layout defining the group. + bias: Shared projection-bias policy defining the group. + """ + _benchmark_multi_head_attention( + benchmark, + case, + attention_type=AttentionType.SELF_ATTENTION, + qk_norm_scope=qk_norm_scope, + rope_interleaved=rope_interleaved, + bias=bias, + ) + + +@pytest.mark.parametrize( + "case", + _IMPLEMENTATION_CASES, + ids=lambda case: case.id, +) +@pytest.mark.parametrize( + "qk_norm_scope,rope_interleaved,bias", + _SHARED_CONFIGS, +) +def test_cross_attention_benchmark( + benchmark: BenchmarkFixture, + case: _ImplementationCase, + qk_norm_scope: QKNormScope, + rope_interleaved: bool, + bias: bool, +) -> None: + """Benchmark end-to-end cross-attention within one shared-policy group. + + Args: + benchmark: Pytest benchmark fixture used to record synchronized timings. + case: Triton backend, precision, and fusion row. + qk_norm_scope: Shared Q/K normalization policy defining the group. + rope_interleaved: Shared rotary-pair layout defining the group. + bias: Shared projection-bias policy defining the group. + """ + _benchmark_multi_head_attention( + benchmark, + case, + attention_type=AttentionType.CROSS_ATTENTION, + qk_norm_scope=qk_norm_scope, + rope_interleaved=rope_interleaved, + bias=bias, + ) diff --git a/flashdreams/benchmarks/recipes/test_wan_modules.py b/flashdreams/benchmarks/recipes/test_wan_modules.py new file mode 100644 index 000000000..ab726fa40 --- /dev/null +++ b/flashdreams/benchmarks/recipes/test_wan_modules.py @@ -0,0 +1,461 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Microbenchmarks for Wan self-attention and DiT blocks. + +Run the manual GPU benchmarks with:: + + uv run --package flashdreams --group test pytest \ + flashdreams/benchmarks/recipes/test_wan_modules.py \ + -p no:manual_marker -m manual --benchmark-only -v +""" + +from __future__ import annotations + +from enum import Enum + +import pytest +import torch +from pytest_benchmark.fixture import BenchmarkFixture + +from flashdreams.accelerated.multi_head_attention_triton import SDPABackend +from flashdreams.core.attention.rope import RotaryPositionEmbedding3D +from flashdreams.recipes.wan.transformer.impl.modules import ( + AttentionBackend, + Block, +) +from flashdreams.recipes.wan.transformer.impl.network import ( + WanDiTNetwork1pt3BConfig, +) + +pytestmark = [ + pytest.mark.manual, + pytest.mark.skipif( + not torch.cuda.is_available(), + reason="Wan DiT module benchmarks require CUDA.", + ), +] + +# Causal-Forcing's chunkwise 480x832 Wan 2.1 1.3B geometry. The VAE produces +# 3x60x104 latents per AR chunk; 1x2x2 DiT patches produce 3x30x52 tokens. +_PIXEL_HEIGHT = 480 +_PIXEL_WIDTH = 832 +_LATENT_HEIGHT = 60 +_LATENT_WIDTH = 104 +_CHUNK_SIZE_T = 3 +_ATTENTION_HEIGHT = 30 +_ATTENTION_WIDTH = 52 +_WINDOW_CHUNKS = 7 +_TEXT_TOKENS = 512 +_CHUNK_TOKENS = _CHUNK_SIZE_T * _ATTENTION_HEIGHT * _ATTENTION_WIDTH +_WINDOW_TOKENS = _WINDOW_CHUNKS * _CHUNK_TOKENS +_SINK_TOKENS = 0 +_WARMUP_ROUNDS = 3 +_BENCHMARK_ROUNDS = 20 +_SEED = 42 + + +class _Implementation(str, Enum): + """Wan attention implementations covered by the benchmark.""" + + WAN_TORCH = "wan_torch" + TRITON_CUDNN = "triton_cudnn" + TRITON_FA2 = "triton_fa2" + + @property + def attention_backend(self) -> AttentionBackend: + """Return the DiT attention implementation.""" + if self is self.WAN_TORCH: + return AttentionBackend.WAN + return AttentionBackend.TRITON + + @property + def sdpa_backend(self) -> SDPABackend: + """Return the configured self-attention SDPA implementation.""" + if self is self.TRITON_FA2: + return SDPABackend.TRITON + return SDPABackend.CUDNN + + @property + def self_attention_operator(self) -> str: + """Return the concrete self-attention operator name.""" + if self is self.WAN_TORCH: + return "cudnn" + if self is self.TRITON_CUDNN: + return "torch_cudnn_sdpa" + return "triton_fa2" + + @property + def cross_attention_operator(self) -> str: + """Return the concrete cross-attention operator name.""" + if self is self.WAN_TORCH: + return "cudnn" + return "triton_fa2" + + +assert {case.attention_backend for case in _Implementation} == set(AttentionBackend) +assert { + case.sdpa_backend + for case in _Implementation + if case.attention_backend is AttentionBackend.TRITON +} == set(SDPABackend) + + +def _skip_unsupported_device( + backend: AttentionBackend, + device: torch.device, +) -> None: + """Skip Triton attention on devices without tensor-memory acceleration.""" + if backend is AttentionBackend.TRITON and torch.cuda.get_device_capability( + device + ) < (9, 0): + pytest.skip("Triton attention requires compute capability 9.0 or newer.") + + +def _make_block( + config: WanDiTNetwork1pt3BConfig, + backend: AttentionBackend, +) -> Block: + """Build a backend-selected block with weight-matched random parameters.""" + + def make(selected_backend: AttentionBackend) -> Block: + return Block( + dim=config.dim, + ffn_dim=config.ffn_dim, + num_heads=config.num_heads, + cross_attn_norm=config.cross_attn_norm, + eps=config.eps, + i2v=config.cross_attn_enable_img, + apply_rope_before_kvcache=config.apply_rope_before_kvcache, + cp_method=config.cp_method, + attention_backend=selected_backend, + sdpa_backend=config.sdpa_backend, + ) + + torch.manual_seed(_SEED) + reference = make(AttentionBackend.WAN) + if backend is AttentionBackend.WAN: + return reference + + block = make(backend) + block.load_state_dict(reference.state_dict(), strict=True) + return block + + +@pytest.mark.parametrize( + "implementation", + tuple(_Implementation), + ids=lambda implementation: implementation.value.replace("_", "-"), +) +@torch.inference_mode() +def test_self_attention_benchmark( + benchmark: BenchmarkFixture, + implementation: _Implementation, +) -> None: + """Benchmark Wan self-attention against a full production KV window.""" + if not torch.cuda.is_bf16_supported(): + pytest.skip("Wan self-attention benchmark requires bfloat16 support.") + + device = torch.device("cuda") + backend = implementation.attention_backend + _skip_unsupported_device(backend, device) + dtype = torch.bfloat16 + config = WanDiTNetwork1pt3BConfig(sdpa_backend=implementation.sdpa_backend) + block = _make_block(config, backend) + assert block.attention_backend is backend + assert block.sdpa_backend is implementation.sdpa_backend + attention = block.self_attn + attention.to(device=device, dtype=dtype).eval() + generator = torch.Generator(device=device).manual_seed(_SEED) + + x = torch.randn( + (_CHUNK_TOKENS, config.dim), + generator=generator, + device=device, + dtype=dtype, + ) + cache = attention.allocate_kv_cache( + batch_size=1, + chunk_size=_CHUNK_TOKENS, + window_size=_WINDOW_TOKENS, + sink_size=_SINK_TOKENS, + device=device, + dtype=dtype, + ) + rope = RotaryPositionEmbedding3D( + head_dim=config.dim // config.num_heads, + len_t=_CHUNK_SIZE_T, + len_h=_ATTENTION_HEIGHT, + len_w=_ATTENTION_WIDTH, + interleaved=True, + device=device, + ) + + benchmark_chunk_idx = _WINDOW_CHUNKS + rope_freqs = [ + rope.shift_t(chunk_idx) for chunk_idx in range(benchmark_chunk_idx + 1) + ] + for chunk_idx in range(_WINDOW_CHUNKS): + cache.before_update(chunk_idx) + output = attention(x, cache, rope_freqs[chunk_idx]) + cache.after_update(chunk_idx) + torch.cuda.synchronize(device) + + benchmark.group = "wan-dit-self-attention" + benchmark.extra_info.update( + { + "module": "self_attention", + "model_family": "wan", + "model_variant": "wan2.1-1.3b", + "implementation": implementation.value, + "attention_backend": backend.value, + "sdpa_backend": implementation.sdpa_backend.value, + "self_attention_operator": implementation.self_attention_operator, + "projection_backend": ( + "separate_qkv" + if backend is AttentionBackend.WAN + else "row_scaled_fp8_fused_qkv_output" + ), + "batch_shape": [], + "flattened_batch_size": 1, + "pixel_resolution": [_PIXEL_HEIGHT, _PIXEL_WIDTH], + "latent_shape": [_CHUNK_SIZE_T, _LATENT_HEIGHT, _LATENT_WIDTH], + "attention_grid": [ + _CHUNK_SIZE_T, + _ATTENTION_HEIGHT, + _ATTENTION_WIDTH, + ], + "chunk_tokens": _CHUNK_TOKENS, + "window_chunks": _WINDOW_CHUNKS, + "window_tokens": _WINDOW_TOKENS, + "sink_tokens": _SINK_TOKENS, + "model_channels": config.dim, + "num_heads": config.num_heads, + "head_dim": config.dim // config.num_heads, + "parameter_count": sum( + parameter.numel() for parameter in attention.parameters() + ), + "checkpoint": "random_init_shared_weights", + "dtype": str(dtype), + "cache_dtype": str(cache.dtype), + "cache_state": "full_window", + "cache_prefill_chunks": _WINDOW_CHUNKS, + "benchmark_chunk_idx": benchmark_chunk_idx, + "cache_update_bookkeeping": "excluded_from_timing", + "rope_interleaved": True, + "context_parallel_size": 1, + "compiled": False, + "cuda_graph": False, + "gpu": torch.cuda.get_device_name(device), + "compute_capability": list(torch.cuda.get_device_capability(device)), + "torch": torch.__version__, + "cuda": torch.version.cuda, + "cudnn": torch.backends.cudnn.version(), + "warmup_rounds": _WARMUP_ROUNDS, + "benchmark_rounds": _BENCHMARK_ROUNDS, + "seed": _SEED, + } + ) + + # Roll outside timing, then repeatedly overwrite the same slot to mirror + # multiple denoising evaluations at one autoregressive position. + cache.before_update(benchmark_chunk_idx) + torch.cuda.synchronize(device) + torch.cuda.reset_peak_memory_stats(device) + + def synchronized_forward() -> torch.Tensor: + result = attention(x, cache, rope_freqs[benchmark_chunk_idx]) + torch.cuda.synchronize(device) + return result + + output = benchmark.pedantic( + synchronized_forward, + iterations=1, + rounds=_BENCHMARK_ROUNDS, + warmup_rounds=_WARMUP_ROUNDS, + ) + cache.after_update(benchmark_chunk_idx) + benchmark.extra_info["peak_cuda_memory_bytes"] = torch.cuda.max_memory_allocated( + device + ) + + assert output.shape == x.shape + assert torch.isfinite(output).all() + + +@pytest.mark.parametrize( + "implementation", + tuple(_Implementation), + ids=lambda implementation: implementation.value.replace("_", "-"), +) +@torch.inference_mode() +def test_dit_block_benchmark( + benchmark: BenchmarkFixture, + implementation: _Implementation, +) -> None: + """Benchmark a production-configured Wan DiT block at steady state.""" + if not torch.cuda.is_bf16_supported(): + pytest.skip("Wan DiT block benchmark requires bfloat16 support.") + + device = torch.device("cuda") + backend = implementation.attention_backend + _skip_unsupported_device(backend, device) + dtype = torch.bfloat16 + config = WanDiTNetwork1pt3BConfig(sdpa_backend=implementation.sdpa_backend) + block = _make_block(config, backend).to(device=device, dtype=dtype).eval() + assert block.attention_backend is backend + assert block.sdpa_backend is implementation.sdpa_backend + block.update_parameters_after_loading_checkpoint() + generator = torch.Generator(device=device).manual_seed(_SEED) + + x = torch.randn( + (_CHUNK_TOKENS, config.dim), + generator=generator, + device=device, + dtype=dtype, + ) + modulation = torch.randn( + (6, config.dim), + generator=generator, + device=device, + dtype=dtype, + ) + context = torch.randn( + (_TEXT_TOKENS, config.dim), + generator=generator, + device=device, + dtype=dtype, + ) + cache = block.initialize_cache( + chunk_size=_CHUNK_TOKENS, + window_size=_WINDOW_TOKENS, + sink_size=_SINK_TOKENS, + context_text=context, + ) + rope = RotaryPositionEmbedding3D( + head_dim=config.dim // config.num_heads, + len_t=_CHUNK_SIZE_T, + len_h=_ATTENTION_HEIGHT, + len_w=_ATTENTION_WIDTH, + interleaved=True, + device=device, + ) + + def forward(chunk_idx: int, chunk_rope_freqs: torch.Tensor) -> torch.Tensor: + cache.before_update(chunk_idx) + result = block( + x=x, + e=modulation, + cache=cache, + rope_freqs=chunk_rope_freqs, + ) + cache.after_update(chunk_idx) + return result + + benchmark_chunk_idx = _WINDOW_CHUNKS + rope_freqs = [ + rope.shift_t(chunk_idx) for chunk_idx in range(benchmark_chunk_idx + 1) + ] + for chunk_idx in range(_WINDOW_CHUNKS): + output = forward(chunk_idx, rope_freqs[chunk_idx]) + torch.cuda.synchronize(device) + + benchmark.group = "wan-dit-block" + benchmark.extra_info.update( + { + "module": "Block", + "model_family": "wan", + "model_variant": "wan2.1-1.3b", + "implementation": implementation.value, + "attention_backend": backend.value, + "sdpa_backend": implementation.sdpa_backend.value, + "self_attention_operator": implementation.self_attention_operator, + "cross_attention_operator": implementation.cross_attention_operator, + "projection_backend": ( + "separate_qkv" + if backend is AttentionBackend.WAN + else "row_scaled_fp8_fused_qkv_output" + ), + "batch_shape": [], + "flattened_batch_size": 1, + "pixel_resolution": [_PIXEL_HEIGHT, _PIXEL_WIDTH], + "latent_shape": [_CHUNK_SIZE_T, _LATENT_HEIGHT, _LATENT_WIDTH], + "attention_grid": [ + _CHUNK_SIZE_T, + _ATTENTION_HEIGHT, + _ATTENTION_WIDTH, + ], + "chunk_tokens": _CHUNK_TOKENS, + "window_chunks": _WINDOW_CHUNKS, + "window_tokens": _WINDOW_TOKENS, + "sink_tokens": _SINK_TOKENS, + "text_tokens": _TEXT_TOKENS, + "model_channels": config.dim, + "ffn_channels": config.ffn_dim, + "num_heads": config.num_heads, + "head_dim": config.dim // config.num_heads, + "parameter_count": sum( + parameter.numel() for parameter in block.parameters() + ), + "checkpoint": "random_init_shared_weights", + "dtype": str(dtype), + "self_attention_cache_dtype": str(cache.self_attn.dtype), + "cross_attention_cache_dtype": str(cache.cross_attn.text.dtype), + "cache_state": "full_window_static_text", + "cache_prefill_chunks": _WINDOW_CHUNKS, + "benchmark_chunk_idx": benchmark_chunk_idx, + "cache_update_bookkeeping": "excluded_from_timing", + "rope_interleaved": True, + "context_parallel_size": 1, + "compiled": False, + "cuda_graph": False, + "gpu": torch.cuda.get_device_name(device), + "compute_capability": list(torch.cuda.get_device_capability(device)), + "torch": torch.__version__, + "cuda": torch.version.cuda, + "cudnn": torch.backends.cudnn.version(), + "warmup_rounds": _WARMUP_ROUNDS, + "benchmark_rounds": _BENCHMARK_ROUNDS, + "seed": _SEED, + } + ) + + cache.before_update(benchmark_chunk_idx) + torch.cuda.synchronize(device) + torch.cuda.reset_peak_memory_stats(device) + + def synchronized_forward() -> torch.Tensor: + result = block( + x=x, + e=modulation, + cache=cache, + rope_freqs=rope_freqs[benchmark_chunk_idx], + ) + torch.cuda.synchronize(device) + return result + + output = benchmark.pedantic( + synchronized_forward, + iterations=1, + rounds=_BENCHMARK_ROUNDS, + warmup_rounds=_WARMUP_ROUNDS, + ) + cache.after_update(benchmark_chunk_idx) + benchmark.extra_info["peak_cuda_memory_bytes"] = torch.cuda.max_memory_allocated( + device + ) + + assert output.shape == x.shape + assert torch.isfinite(output).all() diff --git a/flashdreams/benchmarks/recipes/test_wan_network.py b/flashdreams/benchmarks/recipes/test_wan_network.py new file mode 100644 index 000000000..8def18d31 --- /dev/null +++ b/flashdreams/benchmarks/recipes/test_wan_network.py @@ -0,0 +1,342 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Benchmark the complete Wan 2.1 1.3B DiT network. + +Run the manual GPU benchmarks with:: + + uv run --package flashdreams --group test pytest \ + flashdreams/benchmarks/recipes/test_wan_network.py \ + -p no:manual_marker -m manual --benchmark-only -v +""" + +from __future__ import annotations + +import math +from enum import Enum + +import pytest +import torch +from pytest_benchmark.fixture import BenchmarkFixture + +from flashdreams.accelerated.multi_head_attention_triton import SDPABackend +from flashdreams.core.attention.rope import RotaryPositionEmbedding3D +from flashdreams.infra.acceleration import ( + CUDAGraphDispatch, + cuda_graph_capture_ar_index, +) +from flashdreams.infra.compile import compile_module +from flashdreams.recipes.wan.transformer.impl.modules import AttentionBackend +from flashdreams.recipes.wan.transformer.impl.network import ( + WanDiTNetwork, + WanDiTNetwork1pt3BConfig, +) + +pytestmark = [ + pytest.mark.manual, + pytest.mark.skipif( + not torch.cuda.is_available(), + reason="Wan DiT network benchmarks require CUDA.", + ), +] + +# Causal-Forcing's production chunkwise Wan 2.1 1.3B geometry. +_PIXEL_HEIGHT = 480 +_PIXEL_WIDTH = 832 +_LATENT_HEIGHT = 60 +_LATENT_WIDTH = 104 +_CHUNK_SIZE_T = 3 +_ATTENTION_HEIGHT = 30 +_ATTENTION_WIDTH = 52 +_WINDOW_SIZE_T = 21 +_SINK_SIZE_T = 0 +_TEXT_TOKENS = 512 +_CHUNK_TOKENS = _CHUNK_SIZE_T * _ATTENTION_HEIGHT * _ATTENTION_WIDTH +_WINDOW_CHUNKS = _WINDOW_SIZE_T // _CHUNK_SIZE_T +_WINDOW_TOKENS = _WINDOW_CHUNKS * _CHUNK_TOKENS +_DIFFUSION_TIMESTEP = 1000.0 +_CUDA_GRAPH_WARMUP_ITERS = 2 +_WARMUP_ROUNDS = 3 +_BENCHMARK_ROUNDS = 20 +_SEED = 42 + + +class _Implementation(str, Enum): + """Wan attention implementations covered by the benchmark.""" + + WAN_TORCH = "wan_torch" + TRITON_CUDNN = "triton_cudnn" + TRITON_FA2 = "triton_fa2" + + @property + def attention_backend(self) -> AttentionBackend: + """Return the DiT attention implementation.""" + if self is self.WAN_TORCH: + return AttentionBackend.WAN + return AttentionBackend.TRITON + + @property + def sdpa_backend(self) -> SDPABackend: + """Return the configured self-attention SDPA implementation.""" + if self is self.TRITON_FA2: + return SDPABackend.TRITON + return SDPABackend.CUDNN + + @property + def self_attention_operator(self) -> str: + """Return the concrete self-attention operator name.""" + if self is self.WAN_TORCH: + return "cudnn" + if self is self.TRITON_CUDNN: + return "torch_cudnn_sdpa" + return "triton_fa2" + + @property + def cross_attention_operator(self) -> str: + """Return the concrete cross-attention operator name.""" + if self is self.WAN_TORCH: + return "cudnn" + return "triton_fa2" + + +assert {case.attention_backend for case in _Implementation} == set(AttentionBackend) +assert { + case.sdpa_backend + for case in _Implementation + if case.attention_backend is AttentionBackend.TRITON +} == set(SDPABackend) + + +def _skip_unsupported_device( + backend: AttentionBackend, + device: torch.device, +) -> None: + """Skip Triton attention on devices without tensor-memory acceleration.""" + if backend is AttentionBackend.TRITON and torch.cuda.get_device_capability( + device + ) < (9, 0): + pytest.skip("Triton attention requires compute capability 9.0 or newer.") + + +@pytest.mark.parametrize( + "implementation", + tuple(_Implementation), + ids=lambda implementation: implementation.value.replace("_", "-"), +) +@torch.inference_mode() +def test_dit_network_benchmark( + benchmark: BenchmarkFixture, + implementation: _Implementation, +) -> None: + """Benchmark a compiled Wan 2.1 1.3B DiT using CUDA-graph replay.""" + if not torch.cuda.is_bf16_supported(): + pytest.skip("Wan DiT network benchmark requires bfloat16 support.") + + device = torch.device("cuda") + backend = implementation.attention_backend + _skip_unsupported_device(backend, device) + dtype = torch.bfloat16 + torch.manual_seed(_SEED) + config = WanDiTNetwork1pt3BConfig( + patch_embedding_type="conv3d", + cp_method="ring", + attention_backend=backend, + sdpa_backend=implementation.sdpa_backend, + ) + + # Allocate the real 1.3B network directly in BF16 on its final device. + previous_dtype = torch.get_default_dtype() + try: + torch.set_default_dtype(dtype) + with torch.device(device): + network = WanDiTNetwork(config) + finally: + torch.set_default_dtype(previous_dtype) + network.eval() + network.update_parameters_after_loading_checkpoint() + network.set_context_parallel_group(None) + parameter_count = sum(parameter.numel() for parameter in network.parameters()) + assert all( + block.attention_backend is backend + and block.sdpa_backend is implementation.sdpa_backend + for block in network.blocks + ) + + generator = torch.Generator(device=device).manual_seed(_SEED) + patch_volume = math.prod(config.patch_size) + x = torch.randn( + (_CHUNK_TOKENS, config.in_dim * patch_volume), + generator=generator, + device=device, + dtype=dtype, + ) + timestep = torch.tensor(_DIFFUSION_TIMESTEP, device=device, dtype=dtype) + text_embeddings = torch.randn( + (_TEXT_TOKENS, config.text_dim), + generator=generator, + device=device, + dtype=dtype, + ) + cache = network.initialize_cache( + chunk_size=_CHUNK_TOKENS, + window_size=_WINDOW_TOKENS, + sink_size=0, + text_embeddings=text_embeddings, + ) + rope = RotaryPositionEmbedding3D( + head_dim=config.dim // config.num_heads, + len_t=_CHUNK_SIZE_T, + len_h=_ATTENTION_HEIGHT, + len_w=_ATTENTION_WIDTH, + interleaved=True, + device=device, + ) + + network = compile_module(network) + capture_chunk_idx = cuda_graph_capture_ar_index( + sink_size_t=_SINK_SIZE_T, + window_size_t=_WINDOW_SIZE_T, + len_t=_CHUNK_SIZE_T, + ) + graph_dispatch = CUDAGraphDispatch( + network, + enabled=True, + capture_ar_idx=capture_chunk_idx, + warmup_iters=_CUDA_GRAPH_WARMUP_ITERS, + ) + + def forward(chunk_idx: int, chunk_rope_freqs: torch.Tensor) -> torch.Tensor: + return graph_dispatch.select(chunk_idx, uncond=False)( + x=x, + timesteps=timestep, + cache=cache, + rope_freqs=chunk_rope_freqs, + current_chunk_idx=chunk_idx, + eager_mode=False, + ) + + benchmark_chunk_idx = capture_chunk_idx + 1 + rope_freqs = [ + rope.shift_t(chunk_idx) for chunk_idx in range(benchmark_chunk_idx + 1) + ] + + # Drain compile/autotune while filling the cache. At the first full-window + # index, finish wrapper warmup and capture before benchmarking graph replay. + for chunk_idx in range(capture_chunk_idx): + cache.before_update(chunk_idx) + output = forward(chunk_idx, rope_freqs[chunk_idx]) + cache.after_update(chunk_idx) + cache.before_update(capture_chunk_idx) + for _ in range(_CUDA_GRAPH_WARMUP_ITERS + 1): + output = forward(capture_chunk_idx, rope_freqs[capture_chunk_idx]) + cache.after_update(capture_chunk_idx) + torch.cuda.synchronize(device) + + benchmark.group = "wan-dit-network" + benchmark.extra_info.update( + { + "network": "WanDiTNetwork1pt3B", + "model_family": "wan", + "model_variant": "wan2.1-1.3b", + "implementation": implementation.value, + "execution_backend": "pytorch", + "attention_backend": backend.value, + "sdpa_backend": implementation.sdpa_backend.value, + "self_attention_operator": implementation.self_attention_operator, + "cross_attention_operator": implementation.cross_attention_operator, + "projection_backend": ( + "separate_qkv" + if backend is AttentionBackend.WAN + else "row_scaled_fp8_fused_qkv_output" + ), + "batch_shape": [], + "flattened_batch_size": 1, + "pixel_resolution": [_PIXEL_HEIGHT, _PIXEL_WIDTH], + "latent_shape": [_CHUNK_SIZE_T, _LATENT_HEIGHT, _LATENT_WIDTH], + "attention_grid": [ + _CHUNK_SIZE_T, + _ATTENTION_HEIGHT, + _ATTENTION_WIDTH, + ], + "chunk_tokens": _CHUNK_TOKENS, + "window_chunks": _WINDOW_CHUNKS, + "window_tokens": _WINDOW_TOKENS, + "sink_tokens": 0, + "text_tokens": _TEXT_TOKENS, + "input_patch_channels": config.in_dim * patch_volume, + "output_patch_channels": config.out_dim * patch_volume, + "model_channels": config.dim, + "ffn_channels": config.ffn_dim, + "num_blocks": config.num_layers, + "num_heads": config.num_heads, + "head_dim": config.dim // config.num_heads, + "parameter_count": parameter_count, + "checkpoint": "random_init_seed_matched", + "dtype": str(dtype), + "self_attention_cache_dtype": str(cache[0].self_attn.dtype), + "cross_attention_cache_dtype": str(cache[0].cross_attn.text.dtype), + "compiled": True, + "compile_mode": "max-autotune-no-cudagraphs", + "cuda_graph": True, + "cuda_graph_warmup_iters": _CUDA_GRAPH_WARMUP_ITERS, + "cuda_graph_capture_chunk_idx": capture_chunk_idx, + "cache_state": "full_window_static_text", + "cache_prefill_chunks": capture_chunk_idx + 1, + "benchmark_chunk_idx": benchmark_chunk_idx, + "cache_update_bookkeeping": "excluded_from_timing", + "diffusion_timestep": _DIFFUSION_TIMESTEP, + "classifier_free_guidance": False, + "rope_interleaved": True, + "context_parallel_size": 1, + "gpu": torch.cuda.get_device_name(device), + "compute_capability": list(torch.cuda.get_device_capability(device)), + "torch": torch.__version__, + "cuda": torch.version.cuda, + "cudnn": torch.backends.cudnn.version(), + "warmup_rounds": _WARMUP_ROUNDS, + "benchmark_rounds": _BENCHMARK_ROUNDS, + "compiler_cache_state": ( + "host-dependent; compile, autotune, and CUDA graph capture " + "excluded from measured rounds" + ), + "seed": _SEED, + } + ) + + # Roll once outside timing; all fixture warmups and measured calls replay + # the captured graph against the same steady-state cache slot. + cache.before_update(benchmark_chunk_idx) + torch.cuda.synchronize(device) + torch.cuda.reset_peak_memory_stats(device) + + def synchronized_forward() -> torch.Tensor: + result = forward(benchmark_chunk_idx, rope_freqs[benchmark_chunk_idx]) + torch.cuda.synchronize(device) + return result + + output = benchmark.pedantic( + synchronized_forward, + iterations=1, + rounds=_BENCHMARK_ROUNDS, + warmup_rounds=_WARMUP_ROUNDS, + ) + cache.after_update(benchmark_chunk_idx) + benchmark.extra_info["peak_cuda_memory_bytes"] = torch.cuda.max_memory_allocated( + device + ) + + expected_output_shape = (_CHUNK_TOKENS, config.out_dim * patch_volume) + assert output.shape == expected_output_shape + assert torch.isfinite(output).all() diff --git a/flashdreams/flashdreams/accelerated/fp8_quantization.py b/flashdreams/flashdreams/accelerated/fp8_quantization.py new file mode 100644 index 000000000..a6a6c1125 --- /dev/null +++ b/flashdreams/flashdreams/accelerated/fp8_quantization.py @@ -0,0 +1,124 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Row-scaled E4M3 FP8 quantization and linear projection helpers.""" + +from __future__ import annotations + +import torch +from torch import Tensor + +from flashdreams.accelerated.triton.fp8_quantization import ( + _FP8_MAX, + _quantize_fp8_rows, +) + + +@torch.no_grad() +def quantize_fp8_weight(weight: Tensor) -> tuple[Tensor, Tensor]: + """Quantize a linear weight with one FP32 scale per output row. + + A weight row maps every input feature into one output feature. For output + row ``o``, this stores + ``scale[o] = max(max(abs(weight[o])) / 448, 1e-12)`` and converts + ``weight[o] / scale[o]`` to E4M3. Per-row scaling therefore preserves the + linear layer's ``[O, I]`` layout while giving every output feature its own + dequantization factor. The returned tensors are detached inference data; + gradients continue to belong to the source parameter. + + Args: + weight: Linear weight with shape ``[O, I]``. + + Returns: + Contiguous E4M3 weight with shape ``[O, I]`` and its FP32 scales with + shape ``[O]``. + """ + # Compute one scale per output feature: ``[O, I] -> [O]``. + weight_float = weight.detach().to(torch.float32) + scale = (weight_float.abs().amax(dim=1) / _FP8_MAX).clamp_min(1e-12) + + # Broadcast ``[O] -> [O, 1]`` to normalize each row before E4M3 conversion. + weight_fp8 = ( + (weight_float / scale[:, None]) + .clamp(-_FP8_MAX, _FP8_MAX) + .to(torch.float8_e4m3fn) + .contiguous() + ) + return weight_fp8, scale.contiguous() + + +def fp8_linear( + x: Tensor, + weight: Tensor, + weight_scale: Tensor, + bias: Tensor | None, + out_dtype: torch.dtype, +) -> Tensor: + """Apply a row-scaled FP8 GEMM and restore the requested activation dtype. + + Leading activation dimensions are flattened into ``R`` rows for the GEMM. + Native-precision inputs are dynamically quantized with one scale per row; + E4M3 inputs are interpreted as already quantized with unit row scales. + Weight scales broadcast over rows, so each output feature is independently + dequantized before the optional bias is applied. + + Args: + x: Native-precision or E4M3 activations with shape ``[..., I]``. + weight: E4M3 weight with shape ``[O, I]`` produced by + :func:`quantize_fp8_weight`. + weight_scale: Per-output-feature FP32 scales with shape ``[O]``. + bias: Optional output bias with shape ``[O]``. + out_dtype: Activation dtype returned to the caller. + + Returns: + Projected activations with shape ``[..., O]``. + """ + # Collapse arbitrary leading dimensions for GEMM: + # ``[..., I] -> [R, I]``, where ``R = prod(x.shape[:-1])``. + input_shape = x.shape + x_2d = x.reshape(-1, input_shape[-1]) + + # Supply activation scales as ``[R, 1]``. Native inputs are dynamically + # quantized per row; pre-quantized E4M3 inputs use an identity scale. + if x_2d.dtype == torch.float8_e4m3fn: + x_fp8 = x_2d + input_scale = torch.ones( + (x_2d.shape[0], 1), + device=x.device, + dtype=torch.float32, + ) + else: + x_fp8, input_scale = _quantize_fp8_rows(x_2d) + + # Multiply ``[R, I] @ [I, O] -> [R, O]``. The row scales ``[R, 1]`` + # and transposed-weight scales ``[1, O]`` broadcast over that output. + scaled_bias = bias.to(torch.bfloat16) if bias is not None else None + output = torch._scaled_mm( + x_fp8, + weight.T, + input_scale, + weight_scale.reshape(1, -1).contiguous(), + bias=scaled_bias, + out_dtype=torch.bfloat16, + use_fast_accum=False, + ) + if out_dtype != torch.bfloat16: + output = output.to(out_dtype) + + # Restore the original leading dimensions: ``[R, O] -> [..., O]``. + return output.reshape(input_shape[:-1] + (weight.shape[0],)) + + +__all__ = ["fp8_linear", "quantize_fp8_weight"] diff --git a/flashdreams/flashdreams/accelerated/multi_head_attention.py b/flashdreams/flashdreams/accelerated/multi_head_attention.py new file mode 100644 index 000000000..8f25c711e --- /dev/null +++ b/flashdreams/flashdreams/accelerated/multi_head_attention.py @@ -0,0 +1,252 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Multi-head attention interface and shared policy enums.""" + +from __future__ import annotations + +import math +from abc import ABC, abstractmethod +from enum import Enum +from typing import Generic, TypeVar + +from torch import Tensor, nn + +KVCacheT = TypeVar("KVCacheT") +"""Backend-owned K/V cache type passed to attention.""" + + +class AttentionType(str, Enum): + """Relationship between query tokens and cached K/V context.""" + + SELF_ATTENTION = "self_attention" + """Update K/V from each query chunk before attention.""" + + CROSS_ATTENTION = "cross_attention" + """Query a precomputed static K/V cache without updating it.""" + + +class QKNormScope(str, Enum): + """Feature scope for query and key normalization.""" + + NONE = "none" + """Skip query and key normalization.""" + + HEAD = "head" + """Normalize each attention head independently.""" + + INNER = "inner" + """Normalize the complete projected inner width.""" + + +class MultiHeadAttention(nn.Module, ABC, Generic[KVCacheT]): + """Generic multi-head attention interface over an implementation-owned cache. + + The complete attention operation is the extension point so implementations + may fuse projection, normalization, RoPE, cache mutation, attention, and + output projection as needed. Streaming self-attention updates a + caller-prepared rolling cache from ``x``; cross-attention reads precomputed + static K/V without changing it. + + Shape descriptions use ``L`` for query tokens, ``S`` for cached context + tokens, ``H`` for attention heads, and ``D`` for each head's feature + dimension. Leading ``...`` dimensions describe batch or grouping geometry; + query and context layouts may differ when implementations flatten them to + the same batch size. + """ + + query_dim: int + """Input and output token width in ``[..., L, query_dim]`` tensors.""" + + context_dim: int + """Context token width projected into ``H * D`` key and value features.""" + + attention_type: AttentionType + """Whether forward performs self-attention or static cross-attention.""" + + n_heads: int + """Number of query, key, and value heads.""" + + head_dim: int + """Per-head feature dimension ``D``.""" + + inner_dim: int + """Concatenated head width ``H * D``, equal to ``n_heads * head_dim``.""" + + qk_norm_eps: float + """Epsilon used by query and key RMS normalization when enabled.""" + + qk_norm_scope: QKNormScope + """Feature scope used by query and key normalization.""" + + rope_interleaved: bool + """Whether RoPE rotates adjacent feature pairs instead of half splits.""" + + def __init__( + self, + query_dim: int, + n_heads: int = 8, + head_dim: int = 64, + *, + context_dim: int | None = None, + attention_type: AttentionType = AttentionType.SELF_ATTENTION, + qk_norm_eps: float = 1e-6, + qk_norm_scope: QKNormScope = QKNormScope.HEAD, + rope_interleaved: bool = False, + ) -> None: + """Initialize shared attention geometry and implementation policies. + + Args: + query_dim: Feature dimension of input and output tokens. + n_heads: Number of query, key, and value heads. + head_dim: Feature dimension of each attention head. + context_dim: Feature dimension projected into cached keys and values; + ``None`` uses ``query_dim``. Self-attention requires the two + dimensions to match. + attention_type: Whether :meth:`forward` updates a rolling cache or + queries precomputed static context. + qk_norm_eps: Positive finite epsilon used by Q/K RMS normalization. + qk_norm_scope: Normalize each head, normalize all projected heads + jointly, or disable Q/K normalization. + rope_interleaved: Rotate adjacent feature pairs instead of half splits. + + Raises: + TypeError: An enum policy has the wrong type. + ValueError: A dimension or normalization epsilon is invalid, or + self-attention has different query and context dimensions. + """ + super().__init__() + + context_dim = query_dim if context_dim is None else context_dim + if query_dim <= 0: + raise ValueError(f"query_dim must be positive; got {query_dim}") + if context_dim <= 0: + raise ValueError(f"context_dim must be positive; got {context_dim}") + if n_heads <= 0: + raise ValueError(f"n_heads must be positive; got {n_heads}") + if head_dim <= 0: + raise ValueError(f"head_dim must be positive; got {head_dim}") + if not isinstance(attention_type, AttentionType): + raise TypeError( + f"attention_type must be an AttentionType; got {attention_type!r}" + ) + if attention_type is AttentionType.SELF_ATTENTION and query_dim != context_dim: + raise ValueError( + "self-attention requires query_dim to equal context_dim; " + f"got {query_dim} and {context_dim}" + ) + if not isinstance(qk_norm_scope, QKNormScope): + raise TypeError( + f"qk_norm_scope must be a QKNormScope; got {qk_norm_scope!r}" + ) + if not math.isfinite(qk_norm_eps) or qk_norm_eps <= 0: + raise ValueError( + f"qk_norm_eps must be finite and positive; got {qk_norm_eps}" + ) + + # Store the projection geometry and policies shared by every backend. + self.query_dim = query_dim + self.context_dim = context_dim + self.attention_type = attention_type + self.n_heads = n_heads + self.head_dim = head_dim + self.inner_dim = n_heads * head_dim + self.qk_norm_eps = qk_norm_eps + self.qk_norm_scope = qk_norm_scope + self.rope_interleaved = rope_interleaved + + @property + @abstractmethod + def query_projection(self) -> nn.Linear: + """Return the query projection module. + + This logical accessor does not prescribe the module's registered + attribute name. Model adapters can therefore expose checkpoint-native + names while shared attention implementations consume one interface. + """ + + @property + @abstractmethod + def key_projection(self) -> nn.Linear: + """Return the key projection module.""" + + @property + @abstractmethod + def value_projection(self) -> nn.Linear: + """Return the value projection module.""" + + @property + @abstractmethod + def output_projection(self) -> nn.Linear: + """Return the attention output projection module.""" + + @property + @abstractmethod + def query_norm(self) -> nn.Module: + """Return the query normalization module or identity.""" + + @property + @abstractmethod + def key_norm(self) -> nn.Module: + """Return the key normalization module or identity.""" + + @abstractmethod + def compute_kv( + self, + context: Tensor, + rope_freqs: Tensor | None = None, + ) -> KVCacheT: + """Project context and return a precomputed K/V cache. + + Use this stage to materialize static cross-attention context. The cache + is ready for repeated :meth:`forward` calls; implementations decide its + physical layout, precision, and ownership. + + Args: + context: Key/value source, shape ``[..., S, context_dim]``. + rope_freqs: Optional key positional data for the ``S`` context + tokens; ``None`` leaves keys position-independent. + + Returns: + Precomputed cache containing K/V for all ``S`` context tokens. + """ + + @abstractmethod + def forward( + self, + x: Tensor, + kv_cache: KVCacheT, + rope_freqs: Tensor | None = None, + ) -> Tensor: + """Apply the configured attention type to ``x`` and ``kv_cache``. + + Implementations own the complete operation so fused backends need not + expose independently callable cache-update or cache-query stages. + + Args: + x: Query tokens, shape ``[..., L, query_dim]``. + kv_cache: Streaming cache for self-attention or precomputed static + cache for cross-attention. A streaming cache must already be in + its current-chunk update phase. + rope_freqs: Optional query positional data for ``L`` tokens. For + self-attention, the same data is also applied to current keys; + ``None`` disables positional rotation in both stages. + + Returns: + Attention result with shape ``[..., L, query_dim]``. + """ + + +__all__ = ["AttentionType", "MultiHeadAttention", "QKNormScope"] diff --git a/flashdreams/flashdreams/accelerated/multi_head_attention_triton.py b/flashdreams/flashdreams/accelerated/multi_head_attention_triton.py new file mode 100644 index 000000000..54eec1573 --- /dev/null +++ b/flashdreams/flashdreams/accelerated/multi_head_attention_triton.py @@ -0,0 +1,1206 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Triton-accelerated inference-only multi-head attention.""" + +from __future__ import annotations + +import math +from abc import abstractmethod +from collections.abc import Callable +from enum import Enum + +import torch +import torch.nn.functional as F +from torch import Tensor, nn + +from flashdreams.accelerated.fp8_quantization import fp8_linear, quantize_fp8_weight +from flashdreams.accelerated.multi_head_attention import ( + AttentionType, + MultiHeadAttention, + QKNormScope, +) +from flashdreams.accelerated.triton import ( + flash_attention_2_tma, + fused_rms_rope_kv_cache_update, +) +from flashdreams.core.attention import BlockKVCache + + +class SDPABackend(str, Enum): + """Scaled-dot-product attention implementation.""" + + CUDNN = "cudnn" + """Use PyTorch SDPA forced to the cuDNN backend.""" + + TRITON = "triton" + """Use Triton FlashAttention2 (FA2).""" + + +class QKVFusionOption(str, Enum): + """Projection fusion policy.""" + + NONE = "none" + """Project queries, keys, and values independently.""" + + FULL = "full" + """Use one QKV GEMM; query and context feature widths must match.""" + + FUSE_KV = "fuse_kv" + """Use an independent Q GEMM and one KV GEMM, allowing unequal input widths.""" + + +class _DerivedProjectionWeights(nn.Module): + """Hold execution-ready projection tensors derived from master parameters. + + Registering these tensors as nonpersistent buffers gives PyTorch device and + distributed-module traversal without duplicating them in checkpoints. The + owning attention module rebuilds them after state loading and device/dtype + conversion, which also preserves E4M3 storage instead of converting an old + quantized copy to the requested parameter dtype. + """ + + fused_qkv_weight: Tensor | None + """Full-fusion QKV weight, shape ``[3 * H * D, Q]``.""" + + fused_qkv_bias: Tensor | None + """Full-fusion QKV bias, shape ``[3 * H * D]``.""" + + fused_qkv_weight_scale: Tensor | None + """Full-fusion FP32 scales, shape ``[3 * H * D]``.""" + + fused_kv_weight: Tensor | None + """Fused K/V weight, shape ``[2 * H * D, C]``; a QKV-tail view for ``FULL``.""" + + fused_kv_bias: Tensor | None + """Fused K/V bias, shape ``[2 * H * D]``; a QKV-tail view for ``FULL``.""" + + fused_kv_weight_scale: Tensor | None + """Fused K/V FP32 scales, shape ``[2 * H * D]``; a QKV-tail view for ``FULL``.""" + + q_weight_fp8: Tensor | None + """E4M3 query-projection weight.""" + + q_weight_scale: Tensor | None + """FP32 query weight scales.""" + + k_weight_fp8: Tensor | None + """E4M3 key-projection weight.""" + + k_weight_scale: Tensor | None + """FP32 key weight scales.""" + + v_weight_fp8: Tensor | None + """E4M3 value-projection weight.""" + + v_weight_scale: Tensor | None + """FP32 value weight scales.""" + + output_weight_fp8: Tensor | None + """E4M3 output-projection weight, shape ``[Q, H * D]``.""" + + output_weight_scale: Tensor | None + """FP32 output weight scales, shape ``[Q]``.""" + + def __init__(self) -> None: + """Reserve every derived tensor as a nonpersistent module buffer.""" + super().__init__() + # Registering ``None`` reserves each name in ``Module._buffers``. Later + # tensor assignments therefore remain visible to ``Module.to`` and + # distributed wrappers while ``persistent=False`` excludes them from the + # Torch-compatible state dict. + self.register_buffer("fused_qkv_weight", None, persistent=False) + self.register_buffer("fused_qkv_bias", None, persistent=False) + self.register_buffer("fused_qkv_weight_scale", None, persistent=False) + self.register_buffer("fused_kv_weight", None, persistent=False) + self.register_buffer("fused_kv_bias", None, persistent=False) + self.register_buffer("fused_kv_weight_scale", None, persistent=False) + self.register_buffer("q_weight_fp8", None, persistent=False) + self.register_buffer("q_weight_scale", None, persistent=False) + self.register_buffer("k_weight_fp8", None, persistent=False) + self.register_buffer("k_weight_scale", None, persistent=False) + self.register_buffer("v_weight_fp8", None, persistent=False) + self.register_buffer("v_weight_scale", None, persistent=False) + self.register_buffer("output_weight_fp8", None, persistent=False) + self.register_buffer("output_weight_scale", None, persistent=False) + + +def _cache_write_slice(kv_cache: BlockKVCache) -> tuple[int, int, int]: + """Map the current token chunk onto its physical cache write interval. + + The cache normally consumes the complete ``[B, L, H, D]`` K/V chunk. When a + rolling write would overlap an immutable sink prefix, only the trailing + source tokens that fit after the sink are copied. + + Args: + kv_cache: Block cache prepared for the current chunk. + + Returns: + Source-token offset, destination-cache offset, and token count for the + fused K/V write. + """ + write_start, write_end = kv_cache._current_write_bounds() + read_start = 0 + if ( + kv_cache.sink_size > 0 + and not kv_cache._current_chunk_overlaps_sink() + and write_start < kv_cache.sink_size + ): + # Preserve the sink and align the source tail with the write interval. + # This mirrors BlockKVCache.update without staging processed K/V. + write_start = kv_cache.sink_size + write_length = write_end - write_start + read_start = kv_cache.chunk_size - write_length + return int(read_start), int(write_start), int(write_end - write_start) + + +class TritonMultiHeadAttention(MultiHeadAttention[BlockKVCache]): + """Provide inference-only streaming self- and static cross-attention. + + Shape comments use ``B`` for the product of all leading batch dimensions, + ``L`` for the current query/chunk length, ``S`` for visible cached context, + ``H`` for the number of heads, ``D`` for the head dimension, ``Q`` for + ``query_dim``, and ``C`` for ``context_dim``. + + Full-fusion self-attention produces Q/K/V in one GEMM and keeps processed Q + local to :meth:`forward` while a Triton kernel writes normalized, rotated K/V + directly into cache storage. Cross-attention precomputes K/V independently + and reuses that static cache across forward calls. + + Concrete subclasses own their checkpoint-native projection and + normalization modules and map them to the logical module properties. + Callers own the :class:`BlockKVCache` lifecycle: call ``before_update`` before + streaming self-attention and ``after_update`` once every block has consumed + that chunk. + """ + + @property + @abstractmethod + def query_projection(self) -> nn.Linear: + """Return the query projection module.""" + + @property + @abstractmethod + def key_projection(self) -> nn.Linear: + """Return the key projection module.""" + + @property + @abstractmethod + def value_projection(self) -> nn.Linear: + """Return the value projection module.""" + + @property + @abstractmethod + def output_projection(self) -> nn.Linear: + """Return the attention output projection module.""" + + @property + @abstractmethod + def query_norm(self) -> nn.Module: + """Return the query normalization module or identity.""" + + @property + @abstractmethod + def key_norm(self) -> nn.Module: + """Return the key normalization module or identity.""" + + use_fp8: bool + """Whether projection/output GEMMs and supported attention storage use FP8.""" + + sdpa_backend: SDPABackend + """Scaled-dot-product attention implementation selected at construction.""" + + qkv_fusion_option: QKVFusionOption + """Projection fusion policy selected at construction.""" + + _derived_weights: _DerivedProjectionWeights + """Nonpersistent native and FP8 tensors derived from projection parameters.""" + + def __init__( + self, + query_dim: int, + n_heads: int = 8, + head_dim: int = 64, + *, + context_dim: int | None = None, + attention_type: AttentionType = AttentionType.SELF_ATTENTION, + qkv_fusion_option: QKVFusionOption = QKVFusionOption.FULL, + qk_norm_eps: float = 1e-6, + qk_norm_scope: QKNormScope = QKNormScope.HEAD, + rope_interleaved: bool = False, + use_fp8: bool = False, + sdpa_backend: SDPABackend = SDPABackend.CUDNN, + ) -> None: + """Initialize shared Triton attention geometry and policies. + + Args: + query_dim: Feature dimension of input and output tokens. + n_heads: Number of query, key, and value heads. + head_dim: Feature dimension of each attention head. + context_dim: Feature dimension projected into keys and values. + Defaults to ``query_dim``. + attention_type: Whether forward performs self- or cross-attention. + qkv_fusion_option: Projection fusion policy. Full QKV fusion requires + equal query and context feature dimensions. + qk_norm_eps: Epsilon used by Q/K RMS normalization. + qk_norm_scope: Feature scope used by Q/K RMS normalization, or + :attr:`QKNormScope.NONE` to disable normalization. + rope_interleaved: Rotate adjacent feature pairs instead of half splits. + use_fp8: Use row-scaled E4M3 projection/output GEMMs and E4M3 + attention/cache storage with the Triton backend. cuDNN uses the + native activation dtype for attention and cache storage. + sdpa_backend: Scaled-dot-product attention implementation. + + Raises: + TypeError: A fusion or SDPA backend policy has the wrong type. + ValueError: Full QKV fusion has mismatched input widths, + ``head_dim`` is unsupported, or an FP8 input width is not + aligned to 16 features. + """ + super().__init__( + query_dim=query_dim, + n_heads=n_heads, + head_dim=head_dim, + context_dim=context_dim, + attention_type=attention_type, + qk_norm_eps=qk_norm_eps, + qk_norm_scope=qk_norm_scope, + rope_interleaved=rope_interleaved, + ) + if not isinstance(qkv_fusion_option, QKVFusionOption): + raise TypeError( + "qkv_fusion_option must be a QKVFusionOption; " + f"got {qkv_fusion_option!r}" + ) + if ( + qkv_fusion_option is QKVFusionOption.FULL + and self.query_dim != self.context_dim + ): + raise ValueError( + "full QKV fusion requires query_dim to equal context_dim; " + f"got {self.query_dim} and {self.context_dim}" + ) + if not isinstance(sdpa_backend, SDPABackend): + raise TypeError( + f"sdpa_backend must be an SDPABackend; got {sdpa_backend!r}" + ) + if not (16 <= head_dim <= 256 and head_dim & (head_dim - 1) == 0): + raise ValueError( + "accelerated attention requires a power-of-two head_dim in [16, 256]; " + f"got {head_dim}" + ) + if use_fp8 and (query_dim % 16 != 0 or self.context_dim % 16 != 0): + raise ValueError( + "FP8 projections require query_dim and context_dim to be multiples " + f"of 16; got {query_dim} and {self.context_dim}" + ) + + self.qkv_fusion_option = qkv_fusion_option + self.use_fp8 = use_fp8 + self.sdpa_backend = sdpa_backend + + # ---------------------- Initialization ---------------------- # + + def _initialize_derived_weights(self) -> None: + """Build execution weights after concrete checkpoint fields exist. + + Concrete implementations call this after assigning their checkpoint + fields. The logical accessors are valid before fusion buffers or load + hooks read any projection or normalization module. + """ + self._derived_weights = _DerivedProjectionWeights() + self._refresh_derived_weights() + self.register_load_state_dict_post_hook(self._refresh_derived_weights) + + @torch.no_grad() + def _refresh_derived_weights(self, *args: object) -> None: + """Rebuild execution weights from the registered projection parameters. + + The logical projection accessors remain the source of truth even when a + subclass registers them under model-specific names. + This method materializes only the fusion/precision representation selected + for inference, and runs after strict state loading as well as module moves. + + Args: + args: ``(module, incompatible_keys)`` supplied by PyTorch when invoked + as a load-state-dict post-hook; ignored on both hook and direct calls. + """ + del args + + # Select non-overlapping source matrices so each active group is + # quantized exactly once. FULL derives every projection from QKV; + # FUSE_KV keeps Q separate; NONE keeps Q, K, and V separate. + qkv_weight: Tensor | None = None + qkv_bias: Tensor | None = None + kv_weight: Tensor | None = None + kv_bias: Tensor | None = None + q_weight: Tensor | None = self.query_projection.weight + k_weight: Tensor | None = self.key_projection.weight + v_weight: Tensor | None = self.value_projection.weight + if self.qkv_fusion_option is QKVFusionOption.FULL: + qkv_weight = torch.cat( + ( + self.query_projection.weight, + self.key_projection.weight, + self.value_projection.weight, + ), + dim=0, + ).detach() + q_weight = k_weight = v_weight = None + if self.query_projection.bias is not None: + assert ( + self.key_projection.bias is not None + and self.value_projection.bias is not None + ) + qkv_bias = torch.cat( + ( + self.query_projection.bias, + self.key_projection.bias, + self.value_projection.bias, + ), + dim=0, + ).detach() + elif self.qkv_fusion_option is QKVFusionOption.FUSE_KV: + kv_weight = torch.cat( + (self.key_projection.weight, self.value_projection.weight), dim=0 + ).detach() + k_weight = v_weight = None + if self.key_projection.bias is not None: + assert self.value_projection.bias is not None + kv_bias = torch.cat( + (self.key_projection.bias, self.value_projection.bias), dim=0 + ).detach() + + # Clear every execution buffer before repopulating the active policy. + # Refreshes after state loading or module moves cannot retain stale views. + derived = self._derived_weights + derived.fused_qkv_weight = None + derived.fused_qkv_bias = None + derived.fused_qkv_weight_scale = None + derived.fused_kv_weight = None + derived.fused_kv_bias = None + derived.fused_kv_weight_scale = None + derived.q_weight_fp8 = None + derived.q_weight_scale = None + derived.k_weight_fp8 = None + derived.k_weight_scale = None + derived.v_weight_fp8 = None + derived.v_weight_scale = None + derived.output_weight_fp8 = None + derived.output_weight_scale = None + + if qkv_bias is not None: + derived.fused_qkv_bias = qkv_bias + derived.fused_kv_bias = qkv_bias[self.inner_dim :] + if kv_bias is not None: + derived.fused_kv_bias = kv_bias + + if not self.use_fp8: + if qkv_weight is not None: + native_qkv_weight = qkv_weight.contiguous() + derived.fused_qkv_weight = native_qkv_weight + derived.fused_kv_weight = native_qkv_weight[self.inner_dim :] + if kv_weight is not None: + derived.fused_kv_weight = kv_weight.contiguous() + return + + if qkv_weight is not None: + # Q/K/V and fused K/V share one E4M3 weight and scale allocation. + qkv_fp8, qkv_scale = quantize_fp8_weight(qkv_weight) + derived.fused_qkv_weight = qkv_fp8 + derived.fused_qkv_weight_scale = qkv_scale + derived.fused_kv_weight = qkv_fp8[self.inner_dim :] + derived.fused_kv_weight_scale = qkv_scale[self.inner_dim :] + ( + derived.q_weight_fp8, + derived.k_weight_fp8, + derived.v_weight_fp8, + ) = qkv_fp8.split(self.inner_dim) + ( + derived.q_weight_scale, + derived.k_weight_scale, + derived.v_weight_scale, + ) = qkv_scale.split(self.inner_dim) + + if kv_weight is not None: + # K/V share the fused allocation while Q remains independent. + kv_fp8, kv_scale = quantize_fp8_weight(kv_weight) + derived.fused_kv_weight = kv_fp8 + derived.fused_kv_weight_scale = kv_scale + derived.k_weight_fp8, derived.v_weight_fp8 = kv_fp8.split(self.inner_dim) + derived.k_weight_scale, derived.v_weight_scale = kv_scale.split( + self.inner_dim + ) + + if q_weight is not None: + derived.q_weight_fp8, derived.q_weight_scale = quantize_fp8_weight(q_weight) + if k_weight is not None: + derived.k_weight_fp8, derived.k_weight_scale = quantize_fp8_weight(k_weight) + if v_weight is not None: + derived.v_weight_fp8, derived.v_weight_scale = quantize_fp8_weight(v_weight) + + # Output consumes attention results and always owns its quantized matrix. + ( + derived.output_weight_fp8, + derived.output_weight_scale, + ) = quantize_fp8_weight(self.output_projection.weight) + + def _apply( + self, + fn: Callable[[Tensor], Tensor], + recurse: bool = True, + ) -> TritonMultiHeadAttention: + """Transform parameters and rebuild derived weights on their final device. + + Args: + fn: Tensor transformation applied by :class:`torch.nn.Module`. + recurse: Apply ``fn`` recursively to child modules. + + Returns: + This module with derived projection buffers refreshed. + """ + # Move/cast canonical parameters first, then regenerate from those final + # values. Applying ``fn`` directly to an existing E4M3 buffer would either + # change its dtype or preserve scales computed for stale master weights. + module = super()._apply(fn, recurse=recurse) + self._refresh_derived_weights() + return module + + # ------------------------------------------------------------ # + # Public Methods # + # ------------------------------------------------------------ # + + def allocate_kv_cache( + self, + batch_size: int, + chunk_size: int, + window_size: int, + sink_size: int, + device: torch.device | str, + dtype: torch.dtype, + ) -> BlockKVCache: + """Allocate a rolling cache matching the configured precision policy. + + Args: + batch_size: Flattened batch size ``B``. + chunk_size: Number of current tokens ``L`` written per update. + window_size: Number of rolling context tokens retained after the sink. + sink_size: Number of initial context tokens that are never evicted. + device: Device on which to allocate K/V storage. + dtype: Native activation dtype. The Triton FA2 backend uses E4M3 cache + storage when FP8 is enabled; cuDNN retains this dtype. + + Returns: + Block cache with K/V storage shaped + ``[B, sink_size + window_size, H, D]``. + + Raises: + TypeError: FP8 is enabled with an activation dtype other than FP16 or + BF16. + """ + # Keep ``[B, S, H, D]`` as the public cache shape: ``BlockKVCache`` + # rolls and slices axis 1, and both attention backends accept that logical + # order. + cache_shape = ( + batch_size, + sink_size + window_size, + self.n_heads, + self.head_dim, + ) + if self.use_fp8 and dtype not in (torch.float16, torch.bfloat16): + raise TypeError("FP8 projections require FP16 or BF16 activations") + cache_dtype = ( + torch.float8_e4m3fn + if self.use_fp8 and self.sdpa_backend is SDPABackend.TRITON + else dtype + ) + cache = BlockKVCache( + k_shape=cache_shape, + v_shape=cache_shape, + seq_dim=1, + chunk_size=chunk_size, + window_size=window_size, + sink_size=sink_size, + device=device, + dtype=cache_dtype, + ) + if self.qk_norm_scope is QKNormScope.HEAD: + # ``BlockKVCache`` initially allocates a contiguous ``[B, S, H, D]`` + # tensor, whose physical order makes all ``H * D`` features for one + # token adjacent. Per-head normalization and cuDNN attention instead + # consume one head's complete ``[S, D]`` plane at a time. Because the + # allocation is still empty, reinterpret the same bytes as contiguous + # ``[B, H, S, D]``; ``view`` changes sizes/strides but copies nothing. + storage_shape = ( + batch_size, + self.n_heads, + sink_size + window_size, + self.head_dim, + ) + # Transposing the H/S metadata restores the public ``[B, S, H, D]`` + # shape while retaining physical BHSD order. For example, the final + # strides are ``[H*S*D, D, S*D, 1]``: advancing a token within one + # head moves by ``D`` elements, so ``cache[:, :, h, :]`` is dense. + # A later ``transpose(1, 2)`` recovers physical ``[B, H, S, D]`` order + # for cuDNN without rearranging bytes. The full cache is contiguous; + # a shorter filling-phase prefix retains valid head-major strides. + cache._k = cache._k.view(storage_shape).transpose(1, 2) + cache._v = cache._v.view(storage_shape).transpose(1, 2) + return cache + + @torch.no_grad() + def compute_kv( + self, + context: Tensor, + rope_freqs: Tensor | None = None, + ) -> BlockKVCache: + """Project complete context into a reusable static K/V cache. + + Args: + context: Context tokens shaped ``[..., S, C]``. + rope_freqs: Optional key rotation angles shaped ``[S, 1, 1, D]``. + + Returns: + Filled cache with logical K/V shape ``[B, S, H, D]``. Its sequence + length and window both equal S, so subsequent forward calls read all + context. + """ + key, value = self._project_kv(context) + if rope_freqs is not None: + # Position affects key directions used in Q·K; values are never rotated. + key = self._apply_rope(key, rope_freqs) + + # Convert only after normalization/RoPE, keeping those numerically sensitive + # operations in the native activation dtype. + key = self._attention_storage(key) + value = self._attention_storage(value) + return BlockKVCache.from_tensor(key, value, seq_dim=1) + + @torch.no_grad() + def forward( + self, + x: Tensor, + kv_cache: BlockKVCache, + rope_freqs: Tensor | None = None, + ) -> Tensor: + """Apply self- or cross-attention using the configured cache lifecycle. + + Self-attention updates the prepared rolling cache and computes Q in one + backend-owned branch. Cross-attention computes Q while leaving its + precomputed static cache unchanged. + + Args: + x: Query tokens shaped ``[..., L, Q]``. + kv_cache: Prepared rolling cache for self-attention or precomputed + static cache for cross-attention. + rope_freqs: Optional query rotation angles shaped ``[L, 1, 1, D]``. + Self-attention also applies them to current keys. + + Returns: + Output-projected tokens with the same shape and dtype as ``x``. + """ + if self.attention_type is AttentionType.SELF_ATTENTION: + query = self._update_kv_and_compute_query(x, kv_cache, rope_freqs) + else: + query = self._compute_query(x, rope_freqs) + self._validate_cache(kv_cache, x) + + # ``cached_k/v`` expose only the valid prefix while a rolling cache fills, + # and the complete fixed-size buffer after it reaches steady state. + output = self._attention( + query, + kv_cache.cached_k(), + kv_cache.cached_v(), + ) + sequence_length = x.shape[-2] + output = output.reshape(-1, sequence_length, self.inner_dim) + output = self._project_output(output, x.dtype) + return output.reshape(x.shape[:-2] + (sequence_length, self.query_dim)) + + # ------------------------------------------------------------ # + # Private Method # + # ------------------------------------------------------------ # + + # ------------------ Core Attention Methods ------------------ # + + def _compute_query( + self, + query: Tensor, + rope_freqs: Tensor | None, + ) -> Tensor: + """Project, normalize, and optionally rotate query tokens.""" + query = self._project_query(query) + if rope_freqs is not None: + query = self._apply_rope(query, rope_freqs) + return self._attention_storage(query) + + def _update_kv_and_compute_query( + self, + x: Tensor, + kv_cache: BlockKVCache, + rope_freqs: Tensor | None, + ) -> Tensor: + """Update rolling K/V and return the processed current query.""" + if self.qkv_fusion_option is QKVFusionOption.FULL: + self._validate_fused_update_inputs(x, kv_cache, rope_freqs) + sequence_length = x.shape[-2] + x_flat = x.reshape(-1, sequence_length, self.query_dim) + query, key, value = self._project_qkv(x_flat) + ( + cache_read_start, + cache_write_start, + cache_write_length, + ) = _cache_write_slice(kv_cache) + + if self.qk_norm_scope is QKNormScope.NONE: + if not isinstance(self.query_norm, nn.Identity) or not isinstance( + self.key_norm, nn.Identity + ): + raise RuntimeError( + "Q/K normalization modules must use the same policy" + ) + query_weight: Tensor | None = None + key_weight: Tensor | None = None + else: + if not isinstance(self.query_norm, nn.RMSNorm) or not isinstance( + self.key_norm, nn.RMSNorm + ): + raise RuntimeError( + "Q/K normalization modules must use the same policy" + ) + query_weight = self.query_norm.weight + key_weight = self.key_norm.weight + + # Keep processed Q local while the fused kernel writes K/V directly + # into the current physical cache interval. + return fused_rms_rope_kv_cache_update( + query, + key, + value, + kv_cache._k, + kv_cache._v, + query_weight=query_weight, + key_weight=key_weight, + norm_eps=self.qk_norm_eps, + norm_scope=self.qk_norm_scope, + rope_freqs=rope_freqs, + rope_interleaved=self.rope_interleaved, + cache_read_start=cache_read_start, + cache_write_start=cache_write_start, + cache_write_length=cache_write_length, + ) + + self._validate_tokens(x, self.context_dim, "context") + self._validate_cache(kv_cache, x) + if kv_cache._curr_chunk_idx is None: + raise RuntimeError("call kv_cache.before_update() before attention") + if x.shape[-2] != kv_cache.chunk_size: + raise ValueError( + "context sequence length must equal cache " + f"chunk_size={kv_cache.chunk_size}; got {x.shape[-2]}" + ) + + query = self._compute_query(x, rope_freqs) + key, value = self._project_kv(x) + if rope_freqs is not None: + key = self._apply_rope(key, rope_freqs) + kv_cache.update(self._attention_storage(key), self._attention_storage(value)) + return query + + def _attention(self, query: Tensor, key: Tensor, value: Tensor) -> Tensor: + """Apply the configured non-causal scaled-dot-product attention backend. + + Args: + query: Processed queries with shape ``[B, L, H, D]``. + key: Cached keys with shape ``[B, S, H, D]``. + value: Cached values with shape ``[B, S, H, D]``. + + Returns: + Attention output with shape ``[B, L, H, D]``. + """ + if self.sdpa_backend is SDPABackend.CUDNN: + # The module and Triton kernel use token-major ``[B, L/S, H, D]``. + # PyTorch SDPA instead interprets its two middle axes as ``[H, L/S]``. + # These transposes normally change only shape/stride metadata. A cache + # allocated in physical BHSD order exposes each head's S/D plane + # directly here; selecting its full sequence extent is contiguous. + query = query.transpose(1, 2) + key = key.transpose(1, 2) + value = value.transpose(1, 2) + + # Force cuDNN rather than allowing PyTorch to silently fall back to a + # backend with different performance or supported-layout behavior. + with torch.nn.attention.sdpa_kernel( + torch.nn.attention.SDPBackend.CUDNN_ATTENTION + ): + output = F.scaled_dot_product_attention(query, key, value) + + # Restore the module-wide ``[B, L, H, D]`` contract for head merging. + return output.transpose(1, 2) + + # The TMA wrapper natively consumes and returns token-major BSHD tensors. + return flash_attention_2_tma(query, key, value) + + def _apply_rope(self, x: Tensor, rope_freqs: Tensor) -> Tensor: + """Apply rotary position embeddings to token-major head features. + + Args: + x: Projected Q or K tensor shaped ``[..., L, H, D]``. + rope_freqs: Rotation angles shaped ``[L, 1, 1, D]``. + + Returns: + Rotated tensor with the same shape and dtype as ``x``. + + Raises: + ValueError: ``D`` is odd or the angle tensor has the wrong shape. + RuntimeError: Angles and projected tokens occupy different devices. + """ + if x.shape[-1] % 2 != 0: + raise ValueError(f"RoPE requires an even head_dim; got {x.shape[-1]}") + expected_shape = (x.shape[-3], 1, 1, x.shape[-1]) + if tuple(rope_freqs.shape) != expected_shape: + raise ValueError( + f"rope_freqs must have shape {expected_shape}; " + f"got {tuple(rope_freqs.shape)}" + ) + if rope_freqs.device != x.device: + raise RuntimeError("rope_freqs and tokens must be on the same device") + # Remove the two singleton axes supplied by recipe code, then prepend + # enough singleton batch axes to broadcast ``[L, 1, D]`` over every + # leading batch and all H heads without materializing repeated angles. + freqs = rope_freqs[:, 0, 0, :].reshape( + (1,) * (x.ndim - 3) + (x.shape[-3], 1, x.shape[-1]) + ) + cos_freqs = torch.cos(freqs).to(dtype=x.dtype) + sin_freqs = torch.sin(freqs).to(dtype=x.dtype) + + # Build R(x), the vector rotated 90 degrees inside every feature pair. + # Interleaved RoPE pairs (0,1), (2,3), ...; split-half RoPE pairs feature + # i with i + D/2. ``x*cos(theta) + R(x)*sin(theta)`` then performs all + # independent 2-D rotations in parallel. + if self.rope_interleaved: + rotated = torch.stack((-x[..., 1::2], x[..., 0::2]), dim=-1).flatten(-2) + else: + first, second = x.chunk(2, dim=-1) + rotated = torch.cat((-second, first), dim=-1) + return x * cos_freqs + rotated * sin_freqs + + def _attention_storage(self, x: Tensor) -> Tensor: + """Convert processed Q/K/V to the configured attention storage dtype. + + Args: + x: Native FP16/BF16 projected attention tensor. + + Returns: + E4M3 storage for FP8 Triton FA2, otherwise ``x`` unchanged. + + PyTorch's cuDNN SDPA does not accept FP8 Q/K/V, so ``use_fp8`` affects + its projection GEMMs but leaves attention and cache storage native. + """ + if self.use_fp8 and self.sdpa_backend is SDPABackend.TRITON: + return x.to(torch.float8_e4m3fn) + return x + + # ------------------------ Validation ------------------------ # + + def _validate_tokens(self, x: Tensor, feature_dim: int, name: str) -> None: + """Validate a token tensor before a CUDA projection. + + Args: + x: Query or context tokens shaped ``[..., length, feature_dim]``. + feature_dim: Projection input width required by the module. + name: Argument label included in validation errors. + + Raises: + ValueError: ``x`` lacks sequence/feature axes or has the wrong width. + RuntimeError: ``x`` is not CUDA FP16/BF16 or the GPU predates Hopper. + """ + if x.ndim < 2: + raise ValueError( + f"{name} must have shape [..., L, D]; got {tuple(x.shape)}" + ) + if x.shape[-1] != feature_dim: + raise ValueError( + f"{name} feature width must equal {feature_dim}; got {x.shape[-1]}" + ) + if not x.is_cuda or x.dtype not in (torch.float16, torch.bfloat16): + raise RuntimeError( + "TritonMultiHeadAttention requires CUDA FP16 or BF16 inputs" + ) + if torch.cuda.get_device_capability(x.device)[0] < 9: + raise RuntimeError( + "TritonMultiHeadAttention requires compute capability 9.0 or newer" + ) + + def _validate_cache(self, kv_cache: BlockKVCache, x: Tensor) -> None: + """Validate a cache against query/context tokens. + + Args: + kv_cache: Static or rolling cache with logical shape ``[B, S, H, D]``. + x: Public tokens whose leading dimensions determine flattened batch B. + + Raises: + ValueError: Cache rank, sequence axis, batch, head, or feature shape + does not match this attention module. + RuntimeError: Cache device or storage dtype does not match ``x`` and + the configured backend. + """ + if kv_cache.seq_dim != 1 or kv_cache._k.ndim != 4: + raise ValueError( + "TritonMultiHeadAttention requires a [B, S, H, D] cache with seq_dim=1" + ) + # Public leading dimensions such as batch and video view collapse into + # one B axis before projection; cached K/V must use the same flattening. + expected_shape = (math.prod(x.shape[:-2]), self.n_heads, self.head_dim) + cache_shape = (kv_cache._k.shape[0], kv_cache._k.shape[2], kv_cache._k.shape[3]) + if cache_shape != expected_shape: + raise ValueError( + "cache batch, head, and feature dimensions must equal " + f"{expected_shape}; got {cache_shape}" + ) + if kv_cache._v.shape != kv_cache._k.shape: + raise ValueError("Triton attention requires identical K/V cache shapes") + if kv_cache._k.device != x.device or kv_cache._v.device != x.device: + raise RuntimeError("K/V cache tensors must match the input device") + expected_dtype = ( + torch.float8_e4m3fn + if self.use_fp8 and self.sdpa_backend is SDPABackend.TRITON + else x.dtype + ) + if kv_cache._k.dtype != expected_dtype or kv_cache._v.dtype != expected_dtype: + raise RuntimeError(f"K/V cache tensors must use {expected_dtype}") + + def _validate_fused_update_inputs( + self, + x: Tensor, + kv_cache: BlockKVCache, + rope_freqs: Tensor | None, + ) -> None: + """Validate full-fusion update inputs before cache mutation. + + Args: + x: Current self-attention tokens, shape ``[..., L, Q]``. + kv_cache: Prepared cache with K/V shape ``[B, S, H, D]``. + rope_freqs: Optional current-chunk angles, shape ``[L, 1, 1, D]``. + + Raises: + ValueError: Tensor dimensions or cache layout do not match the module. + RuntimeError: Device, dtype, cache lifecycle, or hardware requirements + are not satisfied. + """ + # Validate the public token shape before flattening leading dimensions. + if x.ndim < 2: + raise ValueError(f"x must have shape [..., L, D]; got {tuple(x.shape)}") + if x.shape[-1] != self.query_dim: + raise ValueError( + f"x feature width must equal query_dim={self.query_dim}; " + f"got {x.shape[-1]}" + ) + + # The accelerated path accepts native FP16/BF16 CUDA inputs; FP8 is + # an internal projection and FA2 attention/cache-storage policy. + if not x.is_cuda or x.dtype not in (torch.float16, torch.bfloat16): + raise RuntimeError( + "TritonMultiHeadAttention requires CUDA FP16 or BF16 inputs" + ) + if torch.cuda.get_device_capability(x.device)[0] < 9: + raise RuntimeError( + "TritonMultiHeadAttention requires compute capability 9.0 or newer" + ) + + # The caller prepares cache write bounds before attention. K/V storage + # keeps logical ``[B, S, H, D]`` axes in one of two supported dense layouts. + if kv_cache._curr_chunk_idx is None: + raise RuntimeError("call kv_cache.before_update() before attention") + if kv_cache.seq_dim != 1 or kv_cache._k.ndim != 4: + raise ValueError( + "TritonMultiHeadAttention requires a [B, S, H, D] cache with seq_dim=1" + ) + if x.shape[-2] != kv_cache.chunk_size: + raise ValueError( + f"x sequence length must equal cache chunk_size={kv_cache.chunk_size}; " + f"got {x.shape[-2]}" + ) + + # Leading input dimensions collapse into the cache's single batch axis: + # ``[..., L, Q] -> [B, L, Q]`` where ``B = prod(x.shape[:-2])``. + batch_size = math.prod(x.shape[:-2]) + expected_cache_shape = (batch_size, self.n_heads, self.head_dim) + cache_shape = ( + kv_cache._k.shape[0], + kv_cache._k.shape[2], + kv_cache._k.shape[3], + ) + if cache_shape != expected_cache_shape: + raise ValueError( + "cache batch, head, and feature dimensions must equal " + f"{expected_cache_shape}; got {cache_shape}" + ) + + # K/V share shape, device, storage precision, and dense layout so one + # fused kernel can write them for the configured attention backend. + if kv_cache._v.shape != kv_cache._k.shape: + raise ValueError("Triton attention requires identical K/V cache shapes") + if kv_cache._k.device != x.device or kv_cache._v.device != x.device: + raise RuntimeError("K/V cache tensors must match the input device") + expected_cache_dtype = ( + torch.float8_e4m3fn + if self.use_fp8 and self.sdpa_backend is SDPABackend.TRITON + else x.dtype + ) + if ( + kv_cache._k.dtype != expected_cache_dtype + or kv_cache._v.dtype != expected_cache_dtype + ): + raise RuntimeError(f"K/V cache tensors must use {expected_cache_dtype}") + # ``is_contiguous`` identifies physical BSHD storage. Transposing S/H and + # checking again identifies the alternate physical BHSD storage while the + # tensors retain logical ``[B, S, H, D]`` shapes. Only HEAD normalization + # supports BHSD because each head's ``[S, D]`` plane must be dense; INNER + # normalization instead needs each token's complete ``H * D`` row dense. + token_major = kv_cache._k.is_contiguous() and kv_cache._v.is_contiguous() + head_major = ( + self.qk_norm_scope is QKNormScope.HEAD + and kv_cache._k.transpose(1, 2).is_contiguous() + and kv_cache._v.transpose(1, 2).is_contiguous() + ) + if not token_major and not head_major: + raise RuntimeError( + "K/V cache storage must be dense token-major, or head-major for " + "head-scoped RMSNorm" + ) + + if rope_freqs is not None: + # RoPE coefficients cover this ``L``-token chunk and broadcast across + # flattened batches and ``H`` heads inside the fused kernel. + expected_rope_shape = (x.shape[-2], 1, 1, self.head_dim) + if tuple(rope_freqs.shape) != expected_rope_shape: + raise ValueError( + f"rope_freqs must have shape {expected_rope_shape}; " + f"got {tuple(rope_freqs.shape)}" + ) + if rope_freqs.device != x.device: + raise RuntimeError("rope_freqs and x must be on the same device") + + # ------------------------ Projection ------------------------ # + + def _project_linear( + self, + x: Tensor, + layer: nn.Linear, + weight_fp8: Tensor | None, + weight_scale: Tensor | None, + ) -> Tensor: + """Apply one native or row-scaled FP8 projection. + + Args: + x: Tokens shaped ``[..., length, layer.in_features]``. + layer: Canonical projection supplying native parameters and bias. + weight_fp8: E4M3 execution weight shaped like ``layer.weight``. + weight_scale: Per-output-row dequantization scales. + + Returns: + Projected tokens in ``x.dtype`` with final width + ``layer.out_features``. + + Raises: + RuntimeError: FP8 execution is selected without a derived weight/scale. + """ + if not self.use_fp8: + return layer(x) + if weight_fp8 is None or weight_scale is None: + raise RuntimeError("FP8 projection weight is not initialized") + return fp8_linear(x, weight_fp8, weight_scale, layer.bias, x.dtype) + + def _project_query(self, query: Tensor) -> Tensor: + """Project and normalize queries in token-major head layout. + + Args: + query: Query tokens shaped ``[..., L, Q]``. + + Returns: + Flattened-batch queries shaped ``[B, L, H, D]``. + """ + self._validate_tokens(query, self.query_dim, "query") + sequence_length = query.shape[-2] + + # The projection emits one ``H * D`` feature vector per token. Splitting + # that final axis exposes heads while ``-1`` folds every public leading + # batch dimension into the kernel's single B axis. + query = self._project_linear( + query, + self.query_projection, + self._derived_weights.q_weight_fp8, + self._derived_weights.q_weight_scale, + ).reshape(-1, sequence_length, self.n_heads, self.head_dim) + + if self.qk_norm_scope is QKNormScope.INNER: + # INNER computes one RMS over all ``H * D`` features of a token. + # Flattening only the last two axes preserves B/L; reshape restores + # the head axis expected by attention. + query = self.query_norm(query.flatten(-2)).reshape(query.shape) + else: + # HEAD leaves D last, so RMSNorm runs independently for each head. + # With NONE, ``q_norm`` is Identity and the same layout passes through. + query = self.query_norm(query) + return query + + def _project_kv(self, context: Tensor) -> tuple[Tensor, Tensor]: + """Project and normalize context keys in token-major head layout. + + Args: + context: Context tokens shaped ``[..., S, C]``. + + Returns: + Flattened-batch key and value tensors shaped ``[B, S, H, D]``. + + Raises: + RuntimeError: The selected fused/FP8 execution weights are unavailable. + """ + self._validate_tokens(context, self.context_dim, "context") + sequence_length = context.shape[-2] + head_shape = (-1, sequence_length, self.n_heads, self.head_dim) + if self.qkv_fusion_option is QKVFusionOption.NONE: + # Independent K and V matrices each produce ``H * D`` features. + key = self._project_linear( + context, + self.key_projection, + self._derived_weights.k_weight_fp8, + self._derived_weights.k_weight_scale, + ).reshape(head_shape) + value = self._project_linear( + context, + self.value_projection, + self._derived_weights.v_weight_fp8, + self._derived_weights.v_weight_scale, + ).reshape(head_shape) + else: + # Both fused policies expose exactly ``[K rows; V rows]`` here. FULL + # stores a view of its QKV tail; FUSE_KV owns the K/V allocation. + fused_weight = self._derived_weights.fused_kv_weight + fused_bias = self._derived_weights.fused_kv_bias + fused_scale = self._derived_weights.fused_kv_weight_scale + + if fused_weight is None: + raise RuntimeError("fused K/V weight is not initialized") + if self.use_fp8: + if fused_scale is None: + raise RuntimeError("FP8 K/V weight scales are not initialized") + projected_kv = fp8_linear( + context, + fused_weight, + fused_scale, + fused_bias, + context.dtype, + ) + else: + projected_kv = F.linear(context, fused_weight, fused_bias) + # The fused output axis is ``[K(H*D), V(H*D)]``. Expose that leading + # K/V selector before H and D, then remove it with zero-copy views. + projected_kv = projected_kv.reshape( + -1, + sequence_length, + 2, + self.n_heads, + self.head_dim, + ) + key, value = projected_kv.unbind(dim=2) + + # Normalize keys because their scale affects Q·K logits; values carry + # payload features and intentionally bypass Q/K normalization. + if self.qk_norm_scope is QKNormScope.INNER: + key = self.key_norm(key.flatten(-2)).reshape(key.shape) + else: + key = self.key_norm(key) + return key, value + + def _project_qkv(self, x: Tensor) -> tuple[Tensor, Tensor, Tensor]: + """Project Q/K/V with one fused native or row-scaled FP8 GEMM. + + Args: + x: Flattened-batch input tokens, shape ``[B, L, Q]``. + + Returns: + Query, key, and value tensors, each shaped ``[B, L, H, D]``. + + Raises: + RuntimeError: A required derived QKV weight or scale is unavailable. + """ + if self._derived_weights.fused_qkv_weight is None: + raise RuntimeError("fused QKV weight is not initialized") + if self.use_fp8: + if self._derived_weights.fused_qkv_weight_scale is None: + raise RuntimeError("FP8 QKV weight scales are not initialized") + # Quantize ``B * L`` activation rows independently, then multiply + # ``[B * L, Q] @ [Q, 3 * H * D]`` and restore ``x.dtype``. + qkv = fp8_linear( + x, + self._derived_weights.fused_qkv_weight, + self._derived_weights.fused_qkv_weight_scale, + self._derived_weights.fused_qkv_bias, + x.dtype, + ) + else: + # ``[B, L, Q] @ [Q, 3 * H * D] -> [B, L, 3 * H * D]``. + qkv = F.linear( + x, + self._derived_weights.fused_qkv_weight, + self._derived_weights.fused_qkv_bias, + ) + + # Split the fused projection axis into Q/K/V, heads, and head features: + # ``[B, L, 3 * H * D] -> [B, L, 3, H, D]``. + qkv = qkv.reshape( + -1, + x.shape[-2], + 3, + self.n_heads, + self.head_dim, + ) + query, key, value = qkv.unbind(dim=2) + return query, key, value + + def _project_output(self, x: Tensor, output_dtype: torch.dtype) -> Tensor: + """Apply the native or row-scaled FP8 output projection. + + Args: + x: Head-concatenated attention output, shape ``[B, L, H * D]``. + output_dtype: Native activation dtype returned to the caller. + + Returns: + Projected tokens with shape ``[B, L, Q]``. + + Raises: + RuntimeError: FP8 is enabled but its derived output weight is missing. + """ + if not self.use_fp8: + # ``[B, L, H * D] @ [H * D, Q] -> [B, L, Q]``. + return self.output_projection(x) + if ( + self._derived_weights.output_weight_fp8 is None + or self._derived_weights.output_weight_scale is None + ): + raise RuntimeError("FP8 output weight is not initialized") + # Quantize ``B * L`` attention rows independently before the scaled GEMM. + return fp8_linear( + x, + self._derived_weights.output_weight_fp8, + self._derived_weights.output_weight_scale, + self.output_projection.bias, + output_dtype, + ) + + +__all__ = ["QKVFusionOption", "SDPABackend", "TritonMultiHeadAttention"] diff --git a/flashdreams/flashdreams/accelerated/triton/__init__.py b/flashdreams/flashdreams/accelerated/triton/__init__.py new file mode 100644 index 000000000..f728358cf --- /dev/null +++ b/flashdreams/flashdreams/accelerated/triton/__init__.py @@ -0,0 +1,32 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Public Triton attention and fused preprocessing kernels.""" + +# Importing the FlashAttention surface also installs the device-workspace +# allocator required by its in-kernel TMA tensor descriptors. +from flashdreams.accelerated.triton.flash_attention import ( + flash_attention_2_tma, + is_tma_flash_attention_supported, +) +from flashdreams.accelerated.triton.rms_rope_kv_cache import ( + fused_rms_rope_kv_cache_update, +) + +__all__ = [ + "flash_attention_2_tma", + "fused_rms_rope_kv_cache_update", + "is_tma_flash_attention_supported", +] diff --git a/flashdreams/flashdreams/accelerated/triton/flash_attention.py b/flashdreams/flashdreams/accelerated/triton/flash_attention.py new file mode 100644 index 000000000..12add1019 --- /dev/null +++ b/flashdreams/flashdreams/accelerated/triton/flash_attention.py @@ -0,0 +1,508 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""TMA-backed Triton FlashAttention2 for projected attention tensors.""" + +from __future__ import annotations + +import math + +import torch +from torch import Tensor + +import triton +import triton.language as tl + + +def _allocate_tma_workspace( + size: int, + alignment: int, + stream: int | None, +) -> Tensor: + """Allocate Triton tensor-descriptor workspace on the active CUDA device. + + Triton invokes this registered allocator for descriptor metadata synthesized + by :func:`triton.language.make_tensor_descriptor` inside a kernel. A PyTorch + byte tensor owns the requested device storage; the CUDA allocator supplies + its alignment and observes the active device and stream. + + Args: + size: Required workspace size in bytes. + alignment: Alignment requested by Triton's allocator protocol; the + PyTorch CUDA allocator provides the actual alignment. + stream: CUDA stream handle; ``None`` denotes the current stream. + + Returns: + Byte tensor with shape ``[size]`` on the active CUDA device. + """ + del alignment, stream + return torch.empty(size, device="cuda", dtype=torch.int8) + + +# In-kernel tensor descriptors need a small device allocation at launch time. +triton.set_allocator(_allocate_tma_workspace) + + +_TMA_ATTENTION_CONFIGS = [ + triton.Config( + {"BLOCK_M": block_m, "BLOCK_N": block_n}, + num_warps=num_warps, + num_stages=num_stages, + ) + for block_m, block_n, num_warps, num_stages in ( + (16, 32, 4, 2), + (32, 32, 4, 2), + (64, 32, 4, 3), + (64, 64, 4, 3), + (64, 64, 8, 3), + (128, 32, 4, 3), + (128, 64, 4, 2), + (128, 64, 4, 3), + (128, 64, 8, 3), + (128, 128, 8, 3), + ) +] +"""Candidate query/key tile geometries for FlashAttention autotuning. + +``BLOCK_M`` controls query rows and the FP32 output-accumulator footprint; +``BLOCK_N`` controls each streamed K/V tile. Warp and stage variants let Triton +balance parallel dot products against descriptor-pipeline resource use.""" + + +def _prune_tma_attention_configs( + configs: list[triton.Config], + named_args: dict[str, object], + **meta: object, +) -> list[triton.Config]: + """Drop tiles that waste work or exceed wide-head shared memory. + + This callback runs before benchmarking so short sequences, wide heads, and + one-byte FP8 storage do not compile configurations whose padded work or + accumulator/pipeline footprint cannot be competitive. + + Args: + configs: Candidate autotuning configurations. + named_args: Runtime arguments containing ``query_length``, + ``key_length``, and the tensor ``element_size`` in bytes. + **meta: Compile-time metadata containing ``HEAD_DIM``. + + Returns: + Configurations whose query and key tiles fit the input geometry. + """ + query_length = named_args["query_length"] + key_length = named_args["key_length"] + element_size = named_args["element_size"] + head_dim = meta["HEAD_DIM"] + assert isinstance(query_length, int) + assert isinstance(key_length, int) + assert isinstance(element_size, int) + assert isinstance(head_dim, int) + # Bound each tile by its sequence axis. Wide ``[D]`` accumulators use at + # most 64 query rows to limit SRAM consumption. Larger K/V tiles and the + # 128x64 two-stage pipeline help 16-bit layouts but waste FP8 resources. + maximum_block_m = min(128, max(16, int(triton.next_power_of_2(query_length)))) + if head_dim > 128: + maximum_block_m = min(maximum_block_m, 64) + maximum_block_n = min( + 128 if element_size == 2 else 64, + max(32, int(triton.next_power_of_2(key_length))), + ) + return [ + config + for config in configs + if config.kwargs["BLOCK_M"] <= maximum_block_m + and config.kwargs["BLOCK_N"] <= maximum_block_n + and not ( + element_size == 1 + and config.kwargs == {"BLOCK_M": 128, "BLOCK_N": 64} + and config.num_stages == 2 + ) + ] + + +# Cache the winning tile by logical geometry, sequence strides, and storage +# width. Pointer values and the numeric softmax scale do not change scheduling, +# so they intentionally do not create new autotuning entries. + + +@triton.autotune( + configs=_TMA_ATTENTION_CONFIGS, + key=[ + "num_heads", + "query_length", + "key_length", + "query_stride_l", + "key_stride_s", + "value_stride_s", + "element_size", + "HEAD_DIM", + ], + prune_configs_by={"early_config_prune": _prune_tma_attention_configs}, + cache_results=True, +) +@triton.jit +def _flash_attention_2_tma_kernel( + query_ptr, + key_ptr, + value_ptr, + output_ptr, + query_stride_b, + query_stride_h, + query_stride_l, + query_stride_d: tl.constexpr, + key_stride_b, + key_stride_h, + key_stride_s, + key_stride_d: tl.constexpr, + value_stride_b, + value_stride_h, + value_stride_s, + value_stride_d: tl.constexpr, + output_stride_b, + output_stride_h, + output_stride_l, + output_stride_d: tl.constexpr, + num_heads: tl.constexpr, + query_length: tl.constexpr, + key_length: tl.constexpr, + element_size, + scale, + HEAD_DIM: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + """Apply tiled non-causal FlashAttention2 with TMA loads and stores. + + Inputs are Q ``[B, L, H, D]`` and K/V ``[B, S, H, D]``. Element strides + describe metadata-only ``[B, H, L|S, D]`` views over that storage. The grid + is ``[ceil_div(L, BLOCK_M), B * H]``. Each program loads one + ``[BLOCK_M, D]`` query tile, streams all ``[BLOCK_N, D]`` K/V tiles, and + produces the matching output tile. Only FP32 online-softmax state and the + output accumulator remain resident; no ``[L, S]`` score matrix is stored. + + Args: + query_ptr: Base pointer for logical queries ``[B, L, H, D]``. + key_ptr: Base pointer for logical keys ``[B, S, H, D]``. + value_ptr: Base pointer for logical values ``[B, S, H, D]``. + output_ptr: Base pointer for logical output ``[B, L, H, D]``. + query_stride_b: Query batch stride in elements. + query_stride_h: Query head stride in elements. + query_stride_l: Query-token stride in elements. + query_stride_d: Query-feature stride in elements. + key_stride_b: Key batch stride in elements. + key_stride_h: Key head stride in elements. + key_stride_s: Key-token stride in elements. + key_stride_d: Key-feature stride in elements. + value_stride_b: Value batch stride in elements. + value_stride_h: Value head stride in elements. + value_stride_s: Value-token stride in elements. + value_stride_d: Value-feature stride in elements. + output_stride_b: Output batch stride in elements. + output_stride_h: Output head stride in elements. + output_stride_l: Output-token stride in elements. + output_stride_d: Output-feature stride in elements. + num_heads: Number of batch/head planes per batch item. + query_length: Logical query-token count ``L``. + key_length: Logical key/value-token count ``S``. + element_size: Query storage width in bytes, used by the autotuning key + and configuration pruning. + scale: Multiplier applied to QK scores before softmax. + HEAD_DIM: Compile-time head width ``D``. + BLOCK_M: Compile-time number of query rows owned by one program. + BLOCK_N: Compile-time number of key/value rows loaded per iteration. + """ + # Decode grid axis 1 into one ``(batch, head)`` plane. Grid axis 0 selects + # the ``[BLOCK_M, D]`` query/output tile within that plane. + query_block = tl.program_id(0) + batch_head = tl.program_id(1) + batch = batch_head // num_heads + head = batch_head % num_heads + + # Offset each base pointer to one batch/head plane. The two-dimensional + # descriptors then traverse only token and feature axes, ``[L|S, D]``. A + # block always spans all ``D`` features; query/output descriptors tile the + # token axis by ``BLOCK_M``, while key/value descriptors use ``BLOCK_N``. + query_base = query_ptr + batch * query_stride_b + head * query_stride_h + key_base = key_ptr + batch * key_stride_b + head * key_stride_h + value_base = value_ptr + batch * value_stride_b + head * value_stride_h + output_base = output_ptr + batch * output_stride_b + head * output_stride_h + query_desc = tl.make_tensor_descriptor( + query_base, + shape=[query_length, HEAD_DIM], + strides=[query_stride_l, query_stride_d], + block_shape=[BLOCK_M, HEAD_DIM], + ) + key_desc = tl.make_tensor_descriptor( + key_base, + shape=[key_length, HEAD_DIM], + strides=[key_stride_s, key_stride_d], + block_shape=[BLOCK_N, HEAD_DIM], + ) + value_desc = tl.make_tensor_descriptor( + value_base, + shape=[key_length, HEAD_DIM], + strides=[value_stride_s, value_stride_d], + block_shape=[BLOCK_N, HEAD_DIM], + ) + output_desc = tl.make_tensor_descriptor( + output_base, + shape=[query_length, HEAD_DIM], + strides=[output_stride_l, output_stride_d], + block_shape=[BLOCK_M, HEAD_DIM], + ) + + # Descriptor boundary handling fills out-of-range rows in the final query + # tile and clips the matching output store, so padded query work never + # reaches logical output storage. + query_start = query_block * BLOCK_M + query = query_desc.load([query_start, 0]) + + # Keep only the FlashAttention2 online-softmax state and the output tile in + # SRAM while K/V tiles stream through TMA. + row_max = tl.full((BLOCK_M,), -float("inf"), tl.float32) + denominator = tl.zeros((BLOCK_M,), tl.float32) + accumulator = tl.zeros((BLOCK_M, HEAD_DIM), tl.float32) + # exp2 is cheaper than exp. log2(e) preserves the requested softmax scale + # while expressing the online recurrence in base two. + qk_scale = scale.to(tl.float32) * 1.4426950408889634 + + for key_start in tl.range(0, key_length, BLOCK_N): + # ``[BLOCK_M, D] @ [D, BLOCK_N] -> [BLOCK_M, BLOCK_N]``. + key = key_desc.load([key_start, 0]) + scores = tl.dot(query, tl.trans(key)) * qk_scale + # TMA fills the final partial key tile with zeros, but a zero QK score + # would still contribute to softmax. Replace those phantom columns with + # negative infinity; their zero probability also makes the padded value + # lanes inert without a separate V mask. + if key_length % BLOCK_N != 0: + key_offsets = tl.arange(0, BLOCK_N) + scores = tl.where( + key_start + key_offsets[None, :] < key_length, scores, -float("inf") + ) + + # Rebase the previous numerator and denominator whenever a new row + # maximum appears. FP32 state keeps long cache windows stable. + tile_max = tl.max(scores, axis=1) + next_row_max = tl.maximum(row_max, tile_max) + correction = tl.exp2(row_max - next_row_max) + probabilities = tl.exp2(scores - next_row_max[:, None]) + denominator = denominator * correction + tl.sum(probabilities, axis=1) + + # Accumulate ``P @ V`` into ``[BLOCK_M, D]`` after rebasing the prior + # numerator to the updated per-row exponent origin. + value = value_desc.load([key_start, 0]) + accumulator *= correction[:, None] + accumulator = tl.dot( + probabilities.to(value.dtype), + value, + accumulator, + ) + row_max = next_row_max + + # Normalize each query row in FP32. The descriptor converts to the output + # storage dtype and clips a final partial query tile while writing logical + # output ``[B, L, H, D]``. + output = accumulator / denominator[:, None] + output_desc.store([query_start, 0], output) + + +def _descriptor_layout_supported(x: Tensor) -> bool: + """Return whether ``x`` satisfies TMA tensor-descriptor stride rules. + + Logical ``[B, L|S, H, D]`` storage is addressed as one + ``[B, H, L|S, D]`` metadata view without copying. Both token-major projected + tensors and head-major cache views are valid when their actual strides meet + the descriptor alignment contract. + + Args: + x: Projected tensor with shape ``[B, L|S, H, D]``. + + Returns: + Whether its ``[B, H, L|S, D]`` element strides are positive, + feature-contiguous, and 16-byte aligned on every outer axis. + """ + element_size = x.element_size() + # Public tensors are [B, L, H, D], but each descriptor traverses an [L, D] + # plane selected by its B/H base pointer. TMA requires contiguous features + # and 16-byte alignment for every outer byte stride. + bhld_strides = (x.stride(0), x.stride(2), x.stride(1), x.stride(3)) + return bhld_strides[-1] == 1 and all( + stride > 0 and stride * element_size % 16 == 0 for stride in bhld_strides[:-1] + ) + + +def is_tma_flash_attention_supported( + query: Tensor, + key: Tensor, + value: Tensor, +) -> bool: + """Return whether projected tensors can use the TMA attention kernel. + + Check shape, placement, storage type, head geometry, device capability, and + descriptor strides without allocating output or launching Triton. + + Args: + query: Query tensor with shape ``[B, L, H, D]``. + key: Key tensor with shape ``[B, S, H, D]``. + value: Value tensor with shape ``[B, S, H, D]``. + + Returns: + Whether Q/K/V have compatible shapes, devices, dtypes, head geometry, + and descriptor strides on a TMA-capable GPU. + """ + if query.ndim != 4 or key.ndim != 4 or value.ndim != 4: + return False + if not query.is_cuda or not key.is_cuda or not value.is_cuda: + return False + if query.device != key.device or query.device != value.device: + return False + if query.dtype != key.dtype or query.dtype != value.dtype: + return False + if query.dtype not in (torch.float16, torch.bfloat16, torch.float8_e4m3fn): + return False + + batch_size, _, num_heads, head_dim = query.shape + if key.shape[0] != batch_size or key.shape[2:] != (num_heads, head_dim): + return False + if value.shape != key.shape: + return False + if not (16 <= head_dim <= 256 and head_dim & (head_dim - 1) == 0): + return False + if torch.cuda.get_device_capability(query.device)[0] < 9: + return False + return all(_descriptor_layout_supported(x) for x in (query, key, value)) + + +def flash_attention_2_tma( + query: Tensor, + key: Tensor, + value: Tensor, + *, + scale: float | None = None, +) -> Tensor: + """Apply non-causal TMA FlashAttention2 to logical Q/K/V tensors. + + Compute ``softmax(scale * Q @ K.T) @ V`` independently for every batch/head + plane, without causal masking or dropout. TMA streams K/V tiles while FP32 + online-softmax state avoids materializing the complete score matrix. FP8 + operands are consumed directly; this interface carries no additional tensor + scale metadata. Empty batch, head, or query axes return an empty output, but + the key/value sequence axis must be positive. + + Args: + query: CUDA query tensor with shape ``[B, L, H, D]`` and FP16, BF16, + or E4M3 storage. + key: Same-device and same-dtype key tensor with shape ``[B, S, H, D]``. + value: Value tensor matching ``key`` exactly. + scale: Multiplier applied to QK scores before softmax; ``None`` uses + ``1 / sqrt(D)``. + + Returns: + Attention result with shape ``[B, L, H, D]`` on the query device and in + the query storage dtype. + + Raises: + ValueError: Q/K/V shapes are incompatible or contain an empty key axis. + RuntimeError: Placement, dtype, head geometry, device capability, or + strides do not satisfy the TMA kernel contract. + """ + if query.ndim != 4 or key.ndim != 4 or value.ndim != 4: + raise ValueError("query, key, and value must have shape [B, L, H, D]") + batch_size, query_length, num_heads, head_dim = query.shape + if key.shape[0] != batch_size or key.shape[2:] != (num_heads, head_dim): + raise ValueError("query and key batch, head, and feature dimensions differ") + if value.shape != key.shape: + raise ValueError("key and value must have identical shapes") + key_length = key.shape[1] + if key_length == 0: + raise ValueError("key and value sequence length must be positive") + if not is_tma_flash_attention_supported(query, key, value): + raise RuntimeError( + "TMA FlashAttention2 requires matching CUDA FP16/BF16/FP8 tensors, " + "compute capability 9.0 or newer, a power-of-two head_dim in " + "[16, 256], and tensor-descriptor-compatible strides" + ) + + # Allocate output ``[B, L, H, D]``; empty outer axes require no launch. + output = torch.empty( + query.shape, + device=query.device, + dtype=query.dtype, + ) + if batch_size == 0 or num_heads == 0 or query_length == 0: + return output + + # Reorder logical ``[B, L, H, D]`` strides from ``(B, L, H, D)`` to the + # per-plane descriptor order ``(B, H, L, D)``. This is metadata only and + # does not transpose or copy input tensors; the output retains the public + # logical order. + query_strides = ( + query.stride(0), + query.stride(2), + query.stride(1), + query.stride(3), + ) + key_strides = (key.stride(0), key.stride(2), key.stride(1), key.stride(3)) + value_strides = ( + value.stride(0), + value.stride(2), + value.stride(1), + value.stride(3), + ) + output_strides = ( + output.stride(0), + output.stride(2), + output.stride(1), + output.stride(3), + ) + + # Autotuning selects the launch shape once per geometry and storage width, + # then reuses it. Grid axes cover query tiles and ``B * H`` planes. + def grid(meta: dict[str, int]) -> tuple[int, int]: + """Build the two-dimensional launch grid for an autotuned query tile. + + Args: + meta: Autotuning metadata containing ``BLOCK_M``. + + Returns: + Query-tile count and flattened batch/head plane count. + """ + return ( + triton.cdiv(query_length, meta["BLOCK_M"]), + batch_size * num_heads, + ) + + # All strides are in elements; ``element_size`` distinguishes storage + # widths in the autotuning cache key. + _flash_attention_2_tma_kernel[grid]( + query, + key, + value, + output, + *query_strides, + *key_strides, + *value_strides, + *output_strides, + num_heads, + query_length, + key_length, + query.element_size(), + 1.0 / math.sqrt(head_dim) if scale is None else scale, + HEAD_DIM=head_dim, + ) + return output + + +__all__ = ["flash_attention_2_tma", "is_tma_flash_attention_supported"] diff --git a/flashdreams/flashdreams/accelerated/triton/fp8_quantization.py b/flashdreams/flashdreams/accelerated/triton/fp8_quantization.py new file mode 100644 index 000000000..17eb13406 --- /dev/null +++ b/flashdreams/flashdreams/accelerated/triton/fp8_quantization.py @@ -0,0 +1,148 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Row-scaled E4M3 FP8 activation quantization with Triton.""" + +from __future__ import annotations + +import torch +from torch import Tensor + +import triton +import triton.language as tl + +_FP8_MAX = 448.0 +"""Finite E4M3 saturation bound used to normalize every activation row before +conversion to ``torch.float8_e4m3fn``.""" + + +@triton.jit +def _quantize_fp8_rows_kernel( + input_ptr, + output_ptr, + scale_ptr, + input_stride_row, + input_stride_column: tl.constexpr, + num_columns: tl.constexpr, + FP8_MAX: tl.constexpr, + MIN_SCALE: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + """Quantize one ``[C]`` activation row and store its FP32 scale. + + Each Triton program handles one row of an input matrix ``[R, C]``. Values + are divided by ``max(max(abs(row)) / FP8_MAX, MIN_SCALE)`` before the output + pointer converts them to E4M3. ``BLOCK_SIZE`` pads the reduction to a power + of two; masked lanes load zero and never write, so they cannot change the + scale or the logical output. The output is contiguous ``[R, C]`` even when + input columns are strided, and ``scale_ptr`` stores one FP32 + dequantization scale per row as ``[R]``. + + Args: + input_ptr: Base pointer for the source matrix ``[R, C]``. + output_ptr: Base pointer for the contiguous E4M3 matrix ``[R, C]``. + scale_ptr: Base pointer for the FP32 row scales ``[R]``. + input_stride_row: Source row stride in elements. + input_stride_column: Source column stride in elements. + num_columns: Logical row width ``C``. + FP8_MAX: Largest finite magnitude used for E4M3 saturation. + MIN_SCALE: Positive scale floor used by zero-valued rows. + BLOCK_SIZE: Power-of-two reduction width greater than or equal to ``C``. + """ + # Select row ``r`` and a power-of-two block of candidate columns ``[C_pad]``. + row = tl.program_id(0) + column_offsets = tl.arange(0, BLOCK_SIZE) + column_mask = column_offsets < num_columns + + # Load ``input[r, :]`` as FP32 and zero masked lanes so they do not affect + # the row-wise absolute maximum. ``values`` has shape ``[C_pad]``. + values = tl.load( + input_ptr + row * input_stride_row + column_offsets * input_stride_column, + mask=column_mask, + other=0.0, + ).to(tl.float32) + + # Reduce ``[C_pad] -> []`` to one positive scale for this row, then map the + # row into the finite E4M3 interval. Explicit NaN propagation keeps invalid + # source values observable instead of silently saturating them. + scale = tl.maximum(tl.max(tl.abs(values), axis=0) / FP8_MAX, MIN_SCALE) + quantized = tl.clamp( + values / scale, + -FP8_MAX, + FP8_MAX, + propagate_nan=tl.PropagateNan.ALL, + ) + + # Store the contiguous E4M3 row ``output[r, :]`` and scalar ``scales[r]``. + # The output address uses the logical width rather than the input stride, + # and the mask prevents padded reduction lanes from reaching memory. + tl.store( + output_ptr + row * num_columns + column_offsets, + quantized, + mask=column_mask, + ) + tl.store(scale_ptr + row, scale) + + +def _quantize_fp8_rows(x: Tensor) -> tuple[Tensor, Tensor]: + """Quantize activation rows with one E4M3 scale per row. + + For row ``r``, the stored scale is + ``max(max(abs(x[r])) / FP8_MAX, 1e-12)`` and the quantized row represents + ``x[r] / scale[r]``. Quantization, scale reduction, and both stores execute + in one Triton program per row. Arbitrary positive input strides are passed + to the kernel, but fresh output storage always uses a contiguous ``[R, C]`` + layout. An empty row axis returns correctly shaped tensors without launching + Triton. + + Args: + x: CUDA activation matrix with shape ``[R, C]``. Rows may be strided, + but ``C`` must be positive. + + Returns: + Contiguous E4M3 activations with shape ``[R, C]`` and FP32 row scales + with shape ``[R, 1]``. + + Raises: + ValueError: ``x`` is not a two-dimensional tensor with a positive width. + """ + if x.ndim != 2 or x.shape[1] == 0: + raise ValueError("x must have shape [rows, columns] with columns > 0") + num_rows, num_columns = x.shape + output = torch.empty(x.shape, device=x.device, dtype=torch.float8_e4m3fn) + scales = torch.empty((num_rows, 1), device=x.device, dtype=torch.float32) + if num_rows == 0: + return output, scales + + # ponytail: one program owns each current <=2048-wide row; use a tiled + # two-pass reduction if future projection widths become materially larger. + # Four warps cover ordinary projection widths; the widest current reduction + # uses eight. One stage is sufficient because a program performs one source + # load followed by an in-register reduction and stores, with no tile loop. + block_size = int(triton.next_power_of_2(num_columns)) + _quantize_fp8_rows_kernel[(num_rows,)]( + x, + output, + scales, + x.stride(0), + x.stride(1), + num_columns, + FP8_MAX=_FP8_MAX, + MIN_SCALE=1e-12, + BLOCK_SIZE=block_size, + num_stages=1, + num_warps=8 if block_size >= 2048 else 4, + ) + return output, scales diff --git a/flashdreams/flashdreams/accelerated/triton/rms_rope_kv_cache.py b/flashdreams/flashdreams/accelerated/triton/rms_rope_kv_cache.py new file mode 100644 index 000000000..3358725e0 --- /dev/null +++ b/flashdreams/flashdreams/accelerated/triton/rms_rope_kv_cache.py @@ -0,0 +1,1019 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Fused Q/K RMS normalization, RoPE, and K/V-cache preprocessing. + +The public wrapper consumes projected Q/K/V tensors with logical layout +``[B, L, H, D]``. It returns every processed query, then maps a selected source +interval on ``L`` into an in-place destination interval on cache axis ``S``. +Keys are normalized and rotated before storage; values are stored unchanged. +Returned Q and cached K/V use the cache dtype, including E4M3 storage for the +FP8 attention path, while RMS reductions and trigonometry use higher precision. + +Three launch topologies match the preprocessing domain. Head scope uses one +program per ``[D]`` head without RoPE and tiles heads when RoPE can share its +trigonometry. Inner and disabled normalization use one program for the packed +``[H * D]`` token width; the disabled scope compiles out the RMS reduction. +Shape comments use ``B`` for batch, ``L`` for current tokens, ``S`` for cache +capacity, ``H`` for heads, and ``D`` for the head dimension. +""" + +from __future__ import annotations + +import torch +from torch import Tensor + +import triton +import triton.language as tl +from flashdreams.accelerated.multi_head_attention import QKNormScope + + +@triton.jit +def _fused_head_rms_rope_kv_cache_kernel( + query_ptr, + key_ptr, + value_ptr, + query_output_ptr, + key_cache_ptr, + value_cache_ptr, + query_weight_ptr, + key_weight_ptr, + rope_freqs_ptr, + query_stride_b, + query_stride_l, + query_stride_h, + query_stride_d, + key_stride_b, + key_stride_l, + key_stride_h, + key_stride_d, + value_stride_b, + value_stride_l, + value_stride_h, + value_stride_d, + query_output_stride_b, + query_output_stride_l, + query_output_stride_h, + query_output_stride_d, + key_cache_stride_b, + key_cache_stride_l, + key_cache_stride_h, + key_cache_stride_d, + value_cache_stride_b, + value_cache_stride_l, + value_cache_stride_h, + value_cache_stride_d, + rope_stride_l, + rope_stride_d, + sequence_length, + num_heads, + cache_read_start, + cache_write_start, + cache_write_length, + EPS: tl.constexpr, + HEAD_DIM: tl.constexpr, + APPLY_NORM: tl.constexpr, + APPLY_ROPE: tl.constexpr, + INTERLEAVED: tl.constexpr, + BLOCK_D: tl.constexpr, +): + """Process one query/key head and write its current cache slice. + + Q/K/V and processed Q have logical shape ``[B, L, H, D]``; caches have + logical shape ``[B, S, H, D]``. All supplied strides are element counts, + so the same pointer arithmetic supports token-major tensors and the + metadata-only head-major cache view accepted by the wrapper. + + The one-dimensional grid contains ``B * L * H`` programs. Each program + reduces and rotates one ``[D]`` vector, stores its processed query, and + conditionally writes its key/value vector when the token belongs to the + selected cache interval. ``BLOCK_D`` is the power-of-two lane count covering + ``D``; lanes beyond ``D`` are masked. ``APPLY_NORM``, ``APPLY_ROPE``, and + ``INTERLEAVED`` are compile-time policies, so disabled work and its pointer + loads are removed from the generated kernel. + """ + # One program owns a (batch, token, head) vector: exactly the reduction + # domain for head-scoped RMSNorm. + program = tl.program_id(0) + head = program % num_heads + token_program = program // num_heads + token = token_program % sequence_length + batch = token_program // sequence_length + + # Mask padded lanes in ``[BLOCK_D]`` while loading one physical ``[D]`` + # head. Explicit element strides permit non-contiguous batch, token, and + # head axes; only the feature axis is required to be contiguous. + dim_offsets = tl.arange(0, BLOCK_D) + dim_mask = dim_offsets < HEAD_DIM + query_base = ( + query_ptr + + batch * query_stride_b + + token * query_stride_l + + head * query_stride_h + ) + key_base = ( + key_ptr + batch * key_stride_b + token * key_stride_l + head * key_stride_h + ) + value_base = ( + value_ptr + + batch * value_stride_b + + token * value_stride_l + + head * value_stride_h + ) + query = tl.load( + query_base + dim_offsets * query_stride_d, + mask=dim_mask, + other=0.0, + ) + key = tl.load( + key_base + dim_offsets * key_stride_d, + mask=dim_mask, + other=0.0, + ) + + if APPLY_NORM: + # Reduce RMS statistics in FP32, then cast back after applying the + # learned ``[D]`` weight so the result matches the projection dtype. + query_scale = tl.rsqrt( + tl.sum(query.to(tl.float32) * query.to(tl.float32), axis=0) / HEAD_DIM + EPS + ) + key_scale = tl.rsqrt( + tl.sum(key.to(tl.float32) * key.to(tl.float32), axis=0) / HEAD_DIM + EPS + ) + query_weight = tl.load( + query_weight_ptr + dim_offsets, + mask=dim_mask, + other=0.0, + ) + key_weight = tl.load( + key_weight_ptr + dim_offsets, + mask=dim_mask, + other=0.0, + ) + query = (query * query_scale * query_weight).to(query.dtype) + key = (key * key_scale * key_weight).to(key.dtype) + + if APPLY_ROPE: + # Map every feature lane to its rotation partner within the same head: + # adjacent pairs for interleaved RoPE, otherwise matching half splits. + if INTERLEAVED: + partner_offsets = tl.where( + dim_offsets % 2 == 0, + dim_offsets + 1, + dim_offsets - 1, + ) + rotation_sign = tl.where(dim_offsets % 2 == 0, -1.0, 1.0) + else: + partner_offsets = tl.where( + dim_offsets < HEAD_DIM // 2, + dim_offsets + HEAD_DIM // 2, + dim_offsets - HEAD_DIM // 2, + ) + rotation_sign = tl.where(dim_offsets < HEAD_DIM // 2, -1.0, 1.0) + + # Reload the unnormalized partner and apply its own affine weight plus + # the head's shared scalar RMS scale. This avoids materializing the + # complete normalized Q/K tensors before rotation. + query_partner = tl.load( + query_base + partner_offsets * query_stride_d, + mask=dim_mask, + other=0.0, + ) + key_partner = tl.load( + key_base + partner_offsets * key_stride_d, + mask=dim_mask, + other=0.0, + ) + if APPLY_NORM: + query_partner_weight = tl.load( + query_weight_ptr + partner_offsets, + mask=dim_mask, + other=0.0, + ) + key_partner_weight = tl.load( + key_weight_ptr + partner_offsets, + mask=dim_mask, + other=0.0, + ) + query_partner = (query_partner * query_scale * query_partner_weight).to( + query_partner.dtype + ) + key_partner = (key_partner * key_scale * key_partner_weight).to( + key_partner.dtype + ) + + # RoPE angles ``[L, D]`` are shared across batch and head axes. + frequencies = tl.load( + rope_freqs_ptr + token * rope_stride_l + dim_offsets * rope_stride_d, + mask=dim_mask, + other=0.0, + ).to(tl.float32) + cos_freqs = tl.cos(frequencies).to(query.dtype) + sin_freqs = tl.sin(frequencies).to(query.dtype) + query = query * cos_freqs + query_partner * sin_freqs * rotation_sign + key = key * cos_freqs + key_partner * sin_freqs * rotation_sign + + # Store processed Q for every token. The destination pointer determines the + # final attention-storage dtype, so this store also performs the E4M3 cast. + query_output_base = ( + query_output_ptr + + batch * query_output_stride_b + + token * query_output_stride_l + + head * query_output_stride_h + ) + tl.store( + query_output_base + dim_offsets * query_output_stride_d, + query, + mask=dim_mask, + ) + + # Translate source token ``t`` to destination + # ``cache_write_start + (t - cache_read_start)``. A nonzero read start occurs + # when an immutable sink clips a rolling write; the mask leaves all cache + # positions outside that interval untouched. + cache_offset = token - cache_read_start + cache_token = cache_write_start + cache_offset + cache_mask = dim_mask & (cache_offset >= 0) & (cache_offset < cache_write_length) + key_cache_base = ( + key_cache_ptr + + batch * key_cache_stride_b + + cache_token * key_cache_stride_l + + head * key_cache_stride_h + ) + value_cache_base = ( + value_cache_ptr + + batch * value_cache_stride_b + + cache_token * value_cache_stride_l + + head * value_cache_stride_h + ) + # K uses its normalized/rotated register value. V is loaded only for tokens + # that will be written and remains otherwise unprocessed. + value = tl.load( + value_base + dim_offsets * value_stride_d, + mask=cache_mask, + other=0.0, + ) + tl.store( + key_cache_base + dim_offsets * key_cache_stride_d, + key, + mask=cache_mask, + ) + tl.store( + value_cache_base + dim_offsets * value_cache_stride_d, + value, + mask=cache_mask, + ) + + +@triton.jit +def _fused_tiled_head_rms_rope_kv_cache_kernel( + query_ptr, + key_ptr, + value_ptr, + query_output_ptr, + key_cache_ptr, + value_cache_ptr, + query_weight_ptr, + key_weight_ptr, + rope_freqs_ptr, + query_stride_b, + query_stride_l, + query_stride_h, + query_stride_d, + key_stride_b, + key_stride_l, + key_stride_h, + key_stride_d, + value_stride_b, + value_stride_l, + value_stride_h, + value_stride_d, + query_output_stride_b, + query_output_stride_l, + query_output_stride_h, + query_output_stride_d, + key_cache_stride_b, + key_cache_stride_l, + key_cache_stride_h, + key_cache_stride_d, + value_cache_stride_b, + value_cache_stride_l, + value_cache_stride_h, + value_cache_stride_d, + rope_stride_l, + rope_stride_d, + sequence_length, + num_heads, + cache_read_start, + cache_write_start, + cache_write_length, + EPS: tl.constexpr, + HEAD_DIM: tl.constexpr, + APPLY_NORM: tl.constexpr, + INTERLEAVED: tl.constexpr, + BLOCK_H: tl.constexpr, +): + """Process a head tile while sharing RoPE coefficients across its heads. + + Q/K/V and processed Q have logical shape ``[B, L, H, D]``; caches have + logical shape ``[B, S, H, D]``. The wrapper launches this kernel only for + head-scoped preprocessing with RoPE and a power-of-two ``D``. Consequently + ``tl.arange(0, HEAD_DIM)`` covers the feature width exactly, while the final + head tile still masks lanes whose head index exceeds ``H``. + + The three-dimensional grid is ``[B, L, ceil_div(H, BLOCK_H)]``. Each program + computes independent ``[D]`` RMS statistics for up to ``BLOCK_H`` heads but + evaluates the token's sine and cosine vectors once for the whole tile. + Explicit element strides support both accepted physical cache layouts. + Processed queries are always stored; K/V writes are masked to the selected + source and destination cache intervals. + """ + # Grid axes directly encode batch, token, and head tile, avoiding integer + # division in the hot tiled-RoPE path. + batch = tl.program_id(0) + token = tl.program_id(1) + head_tile = tl.program_id(2) + + # ``HEAD_DIM`` is exact for this dispatch, so only the potentially partial + # head tile needs masking. ``[:, None]`` broadcasts that mask across D. + head_offsets = head_tile * BLOCK_H + tl.arange(0, BLOCK_H) + head_mask = head_offsets < num_heads + dim_offsets = tl.arange(0, HEAD_DIM) + tensor_mask = head_mask[:, None] + + # Form a ``[BLOCK_H, D]`` address tile from independent head and feature + # strides. This remains valid for token-major Q/K/V and head-major caches. + query_base = query_ptr + batch * query_stride_b + token * query_stride_l + key_base = key_ptr + batch * key_stride_b + token * key_stride_l + value_base = value_ptr + batch * value_stride_b + token * value_stride_l + query_offsets = ( + head_offsets[:, None] * query_stride_h + dim_offsets[None, :] * query_stride_d + ) + key_offsets = ( + head_offsets[:, None] * key_stride_h + dim_offsets[None, :] * key_stride_d + ) + value_offsets = ( + head_offsets[:, None] * value_stride_h + dim_offsets[None, :] * value_stride_d + ) + query = tl.load(query_base + query_offsets, mask=tensor_mask, other=0.0) + key = tl.load(key_base + key_offsets, mask=tensor_mask, other=0.0) + + if APPLY_NORM: + # Axis 1 is D, so every head receives its own FP32 RMS statistic while + # the learned ``[D]`` affine vector broadcasts across the head tile. + query_float = query.to(tl.float32) + key_float = key.to(tl.float32) + query_scale = tl.rsqrt( + tl.sum(query_float * query_float, axis=1) / HEAD_DIM + EPS + ) + key_scale = tl.rsqrt(tl.sum(key_float * key_float, axis=1) / HEAD_DIM + EPS) + query_weight = tl.load(query_weight_ptr + dim_offsets) + key_weight = tl.load(key_weight_ptr + dim_offsets) + query = (query * query_scale[:, None] * query_weight[None, :]).to(query.dtype) + key = (key * key_scale[:, None] * key_weight[None, :]).to(key.dtype) + + # RoPE accepts one angle per feature lane. Materialize trigonometry once per + # token and share it across all heads in this tile. Reshaping exposes either + # adjacent ``(2i, 2i + 1)`` pairs or half-split ``(i, i + D / 2)`` pairs. + frequencies = tl.load( + rope_freqs_ptr + token * rope_stride_l + dim_offsets * rope_stride_d + ).to(tl.float32) + if INTERLEAVED: + frequency_pairs = tl.reshape(frequencies, (HEAD_DIM // 2, 2)) + else: + frequency_pairs = tl.reshape( + frequencies, + (2, HEAD_DIM // 2), + ).permute(1, 0) + frequencies_a, frequencies_b = tl.split(frequency_pairs) + cos_a = tl.cos(frequencies_a).to(query.dtype)[None, :] + sin_a = tl.sin(frequencies_a).to(query.dtype)[None, :] + cos_b = tl.cos(frequencies_b).to(query.dtype)[None, :] + sin_b = tl.sin(frequencies_b).to(query.dtype)[None, :] + + # Apply ``(a, b) -> (a cos - b sin, b cos + a sin)`` without reloading + # partners: both members of every pair are already resident in registers. + if INTERLEAVED: + query_pairs = tl.reshape(query, (BLOCK_H, HEAD_DIM // 2, 2)) + key_pairs = tl.reshape(key, (BLOCK_H, HEAD_DIM // 2, 2)) + query_a, query_b = tl.split(query_pairs) + key_a, key_b = tl.split(key_pairs) + query = tl.reshape( + tl.join( + query_a * cos_a - query_b * sin_a, + query_b * cos_b + query_a * sin_b, + ), + (BLOCK_H, HEAD_DIM), + ) + key = tl.reshape( + tl.join( + key_a * cos_a - key_b * sin_a, + key_b * cos_b + key_a * sin_b, + ), + (BLOCK_H, HEAD_DIM), + ) + else: + query_pairs = tl.reshape( + query, + (BLOCK_H, 2, HEAD_DIM // 2), + ).permute(0, 2, 1) + key_pairs = tl.reshape( + key, + (BLOCK_H, 2, HEAD_DIM // 2), + ).permute(0, 2, 1) + query_a, query_b = tl.split(query_pairs) + key_a, key_b = tl.split(key_pairs) + query = tl.reshape( + tl.join( + query_a * cos_a - query_b * sin_a, + query_b * cos_b + query_a * sin_b, + ).permute(0, 2, 1), + (BLOCK_H, HEAD_DIM), + ) + key = tl.reshape( + tl.join( + key_a * cos_a - key_b * sin_a, + key_b * cos_b + key_a * sin_b, + ).permute(0, 2, 1), + (BLOCK_H, HEAD_DIM), + ) + + # Query storage is independent of cache slicing and may cast to E4M3 when + # ``query_output_ptr`` uses the FP8 cache dtype. + query_output_base = ( + query_output_ptr + batch * query_output_stride_b + token * query_output_stride_l + ) + query_output_offsets = ( + head_offsets[:, None] * query_output_stride_h + + dim_offsets[None, :] * query_output_stride_d + ) + tl.store( + query_output_base + query_output_offsets, + query, + mask=tensor_mask, + ) + + # Apply the same source-to-destination token translation as the generic + # kernel. The scalar token predicate broadcasts across ``[BLOCK_H, D]``. + cache_offset = token - cache_read_start + cache_token = cache_write_start + cache_offset + cache_token_mask = (cache_offset >= 0) & (cache_offset < cache_write_length) + cache_mask = tensor_mask & cache_token_mask + key_cache_base = ( + key_cache_ptr + batch * key_cache_stride_b + cache_token * key_cache_stride_l + ) + value_cache_base = ( + value_cache_ptr + + batch * value_cache_stride_b + + cache_token * value_cache_stride_l + ) + key_cache_offsets = ( + head_offsets[:, None] * key_cache_stride_h + + dim_offsets[None, :] * key_cache_stride_d + ) + value_cache_offsets = ( + head_offsets[:, None] * value_cache_stride_h + + dim_offsets[None, :] * value_cache_stride_d + ) + # Store processed K and untouched V; destination pointer types perform any + # native-to-E4M3 cache conversion. + value = tl.load( + value_base + value_offsets, + mask=cache_mask, + other=0.0, + ) + tl.store( + key_cache_base + key_cache_offsets, + key, + mask=cache_mask, + ) + tl.store( + value_cache_base + value_cache_offsets, + value, + mask=cache_mask, + ) + + +@triton.jit +def _fused_inner_rms_rope_kv_cache_kernel( + query_ptr, + key_ptr, + value_ptr, + query_output_ptr, + key_cache_ptr, + value_cache_ptr, + query_weight_ptr, + key_weight_ptr, + rope_freqs_ptr, + query_stride_b, + query_stride_l, + key_stride_b, + key_stride_l, + value_stride_b, + value_stride_l, + query_output_stride_b, + query_output_stride_l, + key_cache_stride_b, + key_cache_stride_l, + value_cache_stride_b, + value_cache_stride_l, + rope_stride_l, + rope_stride_d, + sequence_length, + cache_read_start, + cache_write_start, + cache_write_length, + EPS: tl.constexpr, + HEAD_DIM: tl.constexpr, + INNER_DIM: tl.constexpr, + APPLY_NORM: tl.constexpr, + APPLY_ROPE: tl.constexpr, + INTERLEAVED: tl.constexpr, + BLOCK_INNER: tl.constexpr, +): + """Process one full-inner-width query/key token and update its cache. + + Q/K/V and processed Q use logical shape ``[B, L, H, D]`` but are traversed + as packed ``[B, L, H * D]`` tensors. The kernel receives only batch/token + strides because ``+ inner_offsets`` relies on the physical ``H`` and ``D`` + axes being contiguous. Caches use token-major ``[B, S, H, D]`` storage with + the same packed inner width. + + The grid has ``B * L`` programs. Each program optionally reduces one + ``[H * D]`` RMS domain, rotates pairs independently within each ``D``-wide + head, stores processed Q, and conditionally writes the selected K/V cache + token. ``BLOCK_INNER`` is the power-of-two lane count covering ``H * D``; + padded lanes are masked. ``QKNormScope.NONE`` uses this topology with + ``APPLY_NORM=False``, which removes the reduction and affine loads at compile + time. + """ + # Inner-scoped RMSNorm couples all heads, so one program owns the complete + # projected width for one (batch, token) pair. + program = tl.program_id(0) + token = program % sequence_length + batch = program // sequence_length + + # Load packed Q/K ``[H * D]``. The wrapper verifies ``stride(H) == D`` and + # ``stride(D) == 1``; padded ``BLOCK_INNER`` lanes are masked. + inner_offsets = tl.arange(0, BLOCK_INNER) + inner_mask = inner_offsets < INNER_DIM + query_base = query_ptr + batch * query_stride_b + token * query_stride_l + key_base = key_ptr + batch * key_stride_b + token * key_stride_l + value_base = value_ptr + batch * value_stride_b + token * value_stride_l + query = tl.load(query_base + inner_offsets, mask=inner_mask, other=0.0) + key = tl.load(key_base + inner_offsets, mask=inner_mask, other=0.0) + + if APPLY_NORM: + # One FP32 RMS reduction spans all heads, followed by learned + # ``[H * D]`` weights for Q and K. + query_scale = tl.rsqrt( + tl.sum(query.to(tl.float32) * query.to(tl.float32), axis=0) / INNER_DIM + + EPS + ) + key_scale = tl.rsqrt( + tl.sum(key.to(tl.float32) * key.to(tl.float32), axis=0) / INNER_DIM + EPS + ) + query_weight = tl.load( + query_weight_ptr + inner_offsets, + mask=inner_mask, + other=0.0, + ) + key_weight = tl.load( + key_weight_ptr + inner_offsets, + mask=inner_mask, + other=0.0, + ) + query = (query * query_scale * query_weight).to(query.dtype) + key = (key * key_scale * key_weight).to(key.dtype) + + if APPLY_ROPE: + # RMS may couple H * D, but RoPE never crosses a head boundary. + # Reconstruct each flattened lane's head base and within-head partner. + head_offsets = (inner_offsets // HEAD_DIM) * HEAD_DIM + dim_offsets = inner_offsets % HEAD_DIM + if INTERLEAVED: + partner_dims = tl.where( + dim_offsets % 2 == 0, + dim_offsets + 1, + dim_offsets - 1, + ) + rotation_sign = tl.where(dim_offsets % 2 == 0, -1.0, 1.0) + else: + partner_dims = tl.where( + dim_offsets < HEAD_DIM // 2, + dim_offsets + HEAD_DIM // 2, + dim_offsets - HEAD_DIM // 2, + ) + rotation_sign = tl.where(dim_offsets < HEAD_DIM // 2, -1.0, 1.0) + partner_offsets = head_offsets + partner_dims + + query_partner = tl.load( + query_base + partner_offsets, + mask=inner_mask, + other=0.0, + ) + key_partner = tl.load( + key_base + partner_offsets, + mask=inner_mask, + other=0.0, + ) + if APPLY_NORM: + query_partner_weight = tl.load( + query_weight_ptr + partner_offsets, + mask=inner_mask, + other=0.0, + ) + key_partner_weight = tl.load( + key_weight_ptr + partner_offsets, + mask=inner_mask, + other=0.0, + ) + query_partner = (query_partner * query_scale * query_partner_weight).to( + query_partner.dtype + ) + key_partner = (key_partner * key_scale * key_partner_weight).to( + key_partner.dtype + ) + + # Convert flattened lanes back to feature indices so ``[L, D]`` RoPE + # angles broadcast identically across every head. + frequencies = tl.load( + rope_freqs_ptr + token * rope_stride_l + dim_offsets * rope_stride_d, + mask=inner_mask, + other=0.0, + ).to(tl.float32) + cos_freqs = tl.cos(frequencies).to(query.dtype) + sin_freqs = tl.sin(frequencies).to(query.dtype) + query = query * cos_freqs + query_partner * sin_freqs * rotation_sign + key = key * cos_freqs + key_partner * sin_freqs * rotation_sign + + # Store every processed query lane; the output pointer casts to the cache + # dtype, including E4M3 for the FP8 attention-storage path. + query_output_base = ( + query_output_ptr + batch * query_output_stride_b + token * query_output_stride_l + ) + tl.store(query_output_base + inner_offsets, query, mask=inner_mask) + + # Map the selected source interval ``[cache_read_start, ...]`` to the + # physical cache interval beginning at ``cache_write_start``. The scalar + # token predicate masks every lane when this program is outside the slice. + cache_offset = token - cache_read_start + cache_token = cache_write_start + cache_offset + cache_mask = inner_mask & (cache_offset >= 0) & (cache_offset < cache_write_length) + key_cache_base = ( + key_cache_ptr + batch * key_cache_stride_b + cache_token * key_cache_stride_l + ) + value_cache_base = ( + value_cache_ptr + + batch * value_cache_stride_b + + cache_token * value_cache_stride_l + ) + # Cache processed K beside unnormalized, unrotated V. Stores cast both to + # the cache pointer dtype. + value = tl.load(value_base + inner_offsets, mask=cache_mask, other=0.0) + tl.store(key_cache_base + inner_offsets, key, mask=cache_mask) + tl.store(value_cache_base + inner_offsets, value, mask=cache_mask) + + +def fused_rms_rope_kv_cache_update( + query: Tensor, + key: Tensor, + value: Tensor, + key_cache: Tensor, + value_cache: Tensor, + *, + query_weight: Tensor | None, + key_weight: Tensor | None, + norm_eps: float, + norm_scope: QKNormScope, + rope_freqs: Tensor | None, + rope_interleaved: bool, + cache_read_start: int, + cache_write_start: int, + cache_write_length: int, +) -> Tensor: + """Normalize and rotate Q/K while writing the current K/V cache slice. + + Normalization precedes RoPE. Processed K and unchanged V from source slice + ``[cache_read_start:cache_read_start + cache_write_length]`` are cast to the + cache dtype and written beginning at ``cache_write_start``. Every token still + produces processed Q, including tokens outside the cache-write slice. + + ``QKNormScope.HEAD`` computes one RMS statistic per ``[D]`` head and uses a + shared ``[D]`` affine weight. ``QKNormScope.INNER`` computes one statistic + over each packed ``[H * D]`` token and uses an ``[H * D]`` affine weight. + ``QKNormScope.NONE`` applies no normalization and requires both weights to be + absent. RoPE always pairs features within a head, regardless of normalization + scope. This inference primitive writes caches in place and provides no + autograd implementation. + + Args: + query: Projected queries with logical shape ``[B, L, H, D]``. The feature + axis must be contiguous; inner and disabled normalization also require + the head axis to be packed with stride ``D``. + key: Projected keys matching ``query`` in shape, dtype, device, and + required physical layout. + value: Projected values matching ``query`` in shape, dtype, device, and + required physical layout. Values are neither normalized nor rotated. + key_cache: Mutable dense storage with logical shape ``[B, S, H, D]``. + Token-major storage is accepted for every scope. Head scope also + accepts a logical view whose ``transpose(1, 2)`` is contiguous, + corresponding to physical ``[B, H, S, D]`` storage. + value_cache: Mutable storage matching ``key_cache`` in shape, dtype, + device, and physical layout. + query_weight: Contiguous RMSNorm query weight with shape ``[D]`` for head + scope or ``[H * D]`` for inner scope; ``None`` for disabled + normalization. + key_weight: RMSNorm key weight with the same shape as ``query_weight``; + ``None`` for disabled normalization. Its device and dtype must match + Q/K. + norm_eps: Epsilon used by Q/K RMS normalization; ignored when disabled. + norm_scope: Normalize each head, normalize the complete projected inner + width, or disable normalization. + rope_freqs: Optional full-width RoPE angles with shape ``[L, 1, 1, D]``. + Angles broadcast across batch and head axes and are converted to FP32 + before evaluating sine and cosine. + rope_interleaved: Pair adjacent features when ``True``; otherwise pair + matching positions in the first and second halves of each head. + Ignored when ``rope_freqs`` is ``None``. + cache_read_start: First token on the ``L`` axis copied into the cache. + cache_write_start: First token on the cache ``S`` axis written. + cache_write_length: Number of consecutive source tokens written. Zero + computes all processed queries without changing either cache. + + Returns: + Contiguous processed queries with shape ``[B, L, H, D]`` and the cache + dtype. Processed K and unchanged V are written in place to their caches. + + Raises: + ValueError: Tensor shapes, cache bounds, RoPE geometry, or normalization + inputs violate the kernel contract. + TypeError: ``norm_scope`` is not a :class:`QKNormScope`. + RuntimeError: Tensor devices, dtypes, strides, or cache storage layouts + are incompatible with the kernels. + """ + # Validate logical Q/K/V ``[B, L, H, D]`` and cache ``[B, S, H, D]`` + # geometry before checking storage properties used by the kernels. + if query.ndim != 4 or key.shape != query.shape or value.shape != query.shape: + raise ValueError( + "query, key, and value must have identical [B, L, H, D] shapes" + ) + if key_cache.ndim != 4 or value_cache.shape != key_cache.shape: + raise ValueError( + "key_cache and value_cache must have identical [B, S, H, D] shapes" + ) + batch_size, sequence_length, num_heads, head_dim = query.shape + if key_cache.shape[0] != batch_size or key_cache.shape[2:] != ( + num_heads, + head_dim, + ): + raise ValueError("cache batch, head, and feature dimensions differ from Q/K/V") + tensors = (query, key, value, key_cache, value_cache) + if not all(x.is_cuda for x in tensors): + raise RuntimeError("fused RMS/RoPE/cache update requires CUDA tensors") + if len({x.device for x in tensors}) != 1: + raise RuntimeError("Q/K/V and cache tensors must share a device") + if key.dtype != query.dtype or value.dtype != query.dtype: + raise RuntimeError("Q/K/V tensors must share a dtype") + if key_cache.dtype != value_cache.dtype: + raise RuntimeError("K/V cache tensors must share a dtype") + if query.dtype not in (torch.float16, torch.bfloat16): + raise RuntimeError("fused RMS/RoPE/cache update requires FP16 or BF16") + if key_cache.dtype not in (query.dtype, torch.float8_e4m3fn): + raise RuntimeError( + "K/V cache storage must match Q/K/V or use torch.float8_e4m3fn" + ) + if not isinstance(norm_scope, QKNormScope): + raise TypeError(f"norm_scope must be a QKNormScope; got {norm_scope!r}") + + # Token-major cache strides are ``[S * H * D, H * D, D, 1]``. Head scope + # additionally accepts a logical ``[B, S, H, D]`` transpose view over dense + # ``[B, H, S, D]`` storage; transposing S/H back must therefore be contiguous. + cache_size = key_cache.shape[1] + token_major = key_cache.is_contiguous() and value_cache.is_contiguous() + head_major = ( + norm_scope is QKNormScope.HEAD + and key_cache.transpose(1, 2).is_contiguous() + and value_cache.transpose(1, 2).is_contiguous() + ) + if not token_major and not head_major: + raise RuntimeError( + "K/V cache storage must be dense token-major, or head-major for " + "head-scoped RMSNorm" + ) + if query.stride(-1) != 1 or key.stride(-1) != 1 or value.stride(-1) != 1: + raise RuntimeError("Q/K/V feature dimensions must be contiguous") + + # The enum is the normalization policy; weight presence must agree with it. + apply_norm = norm_scope is not QKNormScope.NONE + if apply_norm and (query_weight is None or key_weight is None): + raise ValueError("HEAD and INNER normalization require Q/K weights") + if not apply_norm and (query_weight is not None or key_weight is not None): + raise ValueError("NONE normalization requires absent Q/K weights") + inner_dim = num_heads * head_dim + if apply_norm: + assert query_weight is not None and key_weight is not None + expected_weight_size = head_dim if norm_scope is QKNormScope.HEAD else inner_dim + if query_weight.shape != (expected_weight_size,) or key_weight.shape != ( + expected_weight_size, + ): + raise ValueError( + f"RMSNorm weights must have shape ({expected_weight_size},)" + ) + if ( + query_weight.device != query.device + or key_weight.device != query.device + or query_weight.dtype != query.dtype + or key_weight.dtype != query.dtype + ): + raise RuntimeError("RMSNorm weights must match the Q/K device and dtype") + + # RoPE stores one full-width ``[D]`` angle vector per current token; batch + # and head axes broadcast inside both kernels. + if rope_freqs is not None: + expected_rope_shape = (sequence_length, 1, 1, head_dim) + if tuple(rope_freqs.shape) != expected_rope_shape: + raise ValueError( + f"rope_freqs must have shape {expected_rope_shape}; " + f"got {tuple(rope_freqs.shape)}" + ) + if rope_freqs.device != query.device: + raise RuntimeError("rope_freqs must be on the Q/K device") + if head_dim % 2 != 0: + raise ValueError("RoPE requires an even head_dim") + + # The source ``L`` interval and destination ``S`` interval must each fit; + # they may begin at different offsets when a rolling cache preserves sinks. + if not (0 <= cache_read_start <= sequence_length): + raise ValueError("cache_read_start is outside the current sequence") + if not (0 <= cache_write_start <= cache_size): + raise ValueError("cache_write_start is outside the cache") + if cache_write_length < 0: + raise ValueError("cache_write_length must be non-negative") + if cache_read_start + cache_write_length > sequence_length: + raise ValueError("cache source slice exceeds the current sequence") + if cache_write_start + cache_write_length > cache_size: + raise ValueError("cache destination slice exceeds cache storage") + + # Match processed Q to cache storage so attention reads Q/K/V uniformly. + # Triton performs arithmetic in the native Q/K dtype and casts these stores + # to E4M3 only when the cache uses FP8 storage. + query_output = torch.empty( + query.shape, + device=query.device, + dtype=key_cache.dtype, + ) + if batch_size == 0 or sequence_length == 0 or num_heads == 0: + return query_output + + # Triton launch arguments cannot be None. Safe tensor aliases occupy unused + # pointer slots; compile-time APPLY_NORM/APPLY_ROPE remove their loads. + query_weight_ptr = query if query_weight is None else query_weight + key_weight_ptr = key if key_weight is None else key_weight + rope_pointer = query if rope_freqs is None else rope_freqs + # These constexpr values produce separate specialized kernels for each + # normalization, RoPE, and pairing policy instead of runtime branches. + common_meta = { + "EPS": norm_eps, + "HEAD_DIM": head_dim, + "APPLY_NORM": apply_norm, + "APPLY_ROPE": rope_freqs is not None, + "INTERLEAVED": rope_interleaved, + } + if norm_scope is QKNormScope.HEAD: + if rope_freqs is not None and head_dim & (head_dim - 1) == 0: + # Power-of-two D permits exact register reshapes into RoPE pairs. + # Tile heads so one trigonometry pass serves several independent RMS + # domains without introducing padded feature lanes. + block_h = min(16, int(triton.next_power_of_2(num_heads))) + head_tiles = triton.cdiv(num_heads, block_h) + grid = (batch_size, sequence_length, head_tiles) + _fused_tiled_head_rms_rope_kv_cache_kernel[grid]( + query, + key, + value, + query_output, + key_cache, + value_cache, + query_weight_ptr, + key_weight_ptr, + rope_pointer, + *query.stride(), + *key.stride(), + *value.stride(), + *query_output.stride(), + *key_cache.stride(), + *value_cache.stride(), + rope_pointer.stride(0), + rope_pointer.stride(-1), + sequence_length, + num_heads, + cache_read_start, + cache_write_start, + cache_write_length, + EPS=norm_eps, + HEAD_DIM=head_dim, + APPLY_NORM=apply_norm, + INTERLEAVED=rope_interleaved, + BLOCK_H=block_h, + num_warps=4, + num_stages=1, + ) + else: + # Without RoPE there is no trigonometry to share. A non-power-of-two + # D also needs the generic kernel's padded feature mask. + block_d = max(16, int(triton.next_power_of_2(head_dim))) + grid = (batch_size * sequence_length * num_heads,) + _fused_head_rms_rope_kv_cache_kernel[grid]( + query, + key, + value, + query_output, + key_cache, + value_cache, + query_weight_ptr, + key_weight_ptr, + rope_pointer, + *query.stride(), + *key.stride(), + *value.stride(), + *query_output.stride(), + *key_cache.stride(), + *value_cache.stride(), + rope_pointer.stride(0), + rope_pointer.stride(-1), + sequence_length, + num_heads, + cache_read_start, + cache_write_start, + cache_write_length, + BLOCK_D=block_d, + **common_meta, + num_warps=4, + num_stages=1, + ) + else: + # INNER and NONE share the contiguous inner-width topology. NONE sets + # APPLY_NORM=False, so compile-time specialization removes the reduction. + # Flattened loads require H and D to form one contiguous physical span. + if not all( + x.stride(-2) == head_dim and x.stride(-1) == 1 for x in (query, key, value) + ): + raise RuntimeError( + "inner-width preprocessing requires contiguous head and feature dimensions" + ) + # One program per ``(batch, token)`` and one power-of-two block cover the + # packed ``H * D`` domain. Larger widths use more warps to parallelize + # the RMS reduction and vector transforms. + block_inner = max(16, int(triton.next_power_of_2(inner_dim))) + grid = (batch_size * sequence_length,) + _fused_inner_rms_rope_kv_cache_kernel[grid]( + query, + key, + value, + query_output, + key_cache, + value_cache, + query_weight_ptr, + key_weight_ptr, + rope_pointer, + query.stride(0), + query.stride(1), + key.stride(0), + key.stride(1), + value.stride(0), + value.stride(1), + query_output.stride(0), + query_output.stride(1), + key_cache.stride(0), + key_cache.stride(1), + value_cache.stride(0), + value_cache.stride(1), + rope_pointer.stride(0), + rope_pointer.stride(-1), + sequence_length, + cache_read_start, + cache_write_start, + cache_write_length, + INNER_DIM=inner_dim, + BLOCK_INNER=block_inner, + **common_meta, + num_warps=8 if inner_dim > 2048 else 4, + num_stages=1, + ) + return query_output + + +__all__ = ["fused_rms_rope_kv_cache_update"] diff --git a/flashdreams/flashdreams/infra/diffusion/model/base.py b/flashdreams/flashdreams/infra/diffusion/model/base.py index 57b4e5171..4f1e598f0 100644 --- a/flashdreams/flashdreams/infra/diffusion/model/base.py +++ b/flashdreams/flashdreams/infra/diffusion/model/base.py @@ -20,6 +20,7 @@ from dataclasses import dataclass, field from typing import Any, Generic, cast +import nvtx import torch import torch.nn as nn from torch import Tensor @@ -197,12 +198,13 @@ def predict_flow(noisy_latent: Tensor, timestep: Tensor) -> Tensor: noisy_latent ) - output = self.transformer.predict_flow( - noisy_latent=noisy_latent, - timestep=timestep, - cache=cache, - input=input, - ) + with nvtx.annotate("flashdreams.diffusion.predict_flow"): + output = self.transformer.predict_flow( + noisy_latent=noisy_latent, + timestep=timestep, + cache=cache, + input=input, + ) if self.config.noise_in_unpatchified_shape: output = self.transformer.unpatchify_and_maybe_gather_cp(output) diff --git a/flashdreams/flashdreams/infra/pipeline/base.py b/flashdreams/flashdreams/infra/pipeline/base.py index 0dc546927..3eae70041 100644 --- a/flashdreams/flashdreams/infra/pipeline/base.py +++ b/flashdreams/flashdreams/infra/pipeline/base.py @@ -20,6 +20,7 @@ from dataclasses import dataclass, field from typing import Any, Generic +import nvtx import torch import torch.nn as nn from loguru import logger @@ -147,6 +148,7 @@ def __init__(self, config: StreamInferencePipelineConfig) -> None: def device(self) -> torch.device: return self.diffusion_model.device + @nvtx.annotate("flashdreams.pipeline.initialize_cache") def initialize_cache( self, transformer_context: dict[str, Any] | None = None, @@ -191,6 +193,7 @@ def initialize_cache( ) @torch.no_grad() + @nvtx.annotate("flashdreams.pipeline.generate") def generate( self, autoregressive_index: int, @@ -234,20 +237,22 @@ def generate( "NullEncoderConfig() for an identity passthrough)." ) assert cache.encoder_cache is not None # invariant: paired with encoder - input = self.encoder( - input=input, - autoregressive_index=autoregressive_index, - cache=cache.encoder_cache, - ) + with nvtx.annotate("flashdreams.pipeline.encode"): + input = self.encoder( + input=input, + autoregressive_index=autoregressive_index, + cache=cache.encoder_cache, + ) if events is not None: events.record("encode") - clean_latent, final_state = self.diffusion_model.generate( - autoregressive_index=autoregressive_index, - cache=cache.transformer_cache, - input=input, - ) + with nvtx.annotate("flashdreams.pipeline.diffuse"): + clean_latent, final_state = self.diffusion_model.generate( + autoregressive_index=autoregressive_index, + cache=cache.transformer_cache, + input=input, + ) cache.final_state = final_state if events is not None: @@ -255,11 +260,12 @@ def generate( if self.decoder is not None: assert cache.decoder_cache is not None # invariant: paired with decoder - output = self.decoder( - input=clean_latent, - autoregressive_index=autoregressive_index, - cache=cache.decoder_cache, - ) + with nvtx.annotate("flashdreams.pipeline.decode"): + output = self.decoder( + input=clean_latent, + autoregressive_index=autoregressive_index, + cache=cache.decoder_cache, + ) else: output = clean_latent @@ -269,6 +275,7 @@ def generate( return output @torch.no_grad() + @nvtx.annotate("flashdreams.pipeline.finalize") def finalize( self, autoregressive_index: int, diff --git a/flashdreams/flashdreams/recipes/cosmos/transformer/impl/modules.py b/flashdreams/flashdreams/recipes/cosmos/transformer/impl/modules.py index 5144c9918..cb27df065 100644 --- a/flashdreams/flashdreams/recipes/cosmos/transformer/impl/modules.py +++ b/flashdreams/flashdreams/recipes/cosmos/transformer/impl/modules.py @@ -329,7 +329,7 @@ def update_kv( """Append K/V computed from ``x`` into an existing ``kv_cache``.""" return self._compute_or_update_kv_cache(x, kv_cache, rope_freqs) - def apply_kv( + def query_kv( self, x: Tensor, kv_cache: BlockKVCache, @@ -382,13 +382,13 @@ def forward( """ if update_kv_cache: kv_cache = self.update_kv(x, kv_cache, rope_freqs) - return self.apply_kv(x, kv_cache, rope_freqs) + return self.query_kv(x, kv_cache, rope_freqs) class SelfAttention(MultiHeadAttention): """Self-attention: queries and K/V are derived from the same ``x`` each step.""" - def initialize_cache( + def allocate_kv_cache( self, batch_size: int, chunk_size: int, @@ -397,7 +397,7 @@ def initialize_cache( device: torch.device, dtype: torch.dtype, ) -> BlockKVCache: - """Initialize KV cache for streaming self-attention. + """Allocate a KV cache for streaming self-attention. Args: batch_size: Flattened batch size used by attention. @@ -565,7 +565,7 @@ def initialize_cache( batch_shape = context.shape[:-2] batch_size = math.prod(batch_shape) return BlockCache( - self_attn=self.self_attn.initialize_cache( + self_attn=self.self_attn.allocate_kv_cache( batch_size, chunk_size, window_size, diff --git a/flashdreams/flashdreams/recipes/wan/transformer/impl/modules.py b/flashdreams/flashdreams/recipes/wan/transformer/impl/modules.py index 5c68e4c28..4b1862730 100644 --- a/flashdreams/flashdreams/recipes/wan/transformer/impl/modules.py +++ b/flashdreams/flashdreams/recipes/wan/transformer/impl/modules.py @@ -19,6 +19,7 @@ import math from dataclasses import dataclass +from enum import Enum from typing import Any, Literal import torch @@ -26,6 +27,15 @@ from torch import Tensor from torch.distributed import ProcessGroup +from flashdreams.accelerated.multi_head_attention import ( + AttentionType, + QKNormScope, +) +from flashdreams.accelerated.multi_head_attention_triton import ( + QKVFusionOption, + SDPABackend, + TritonMultiHeadAttention, +) from flashdreams.core.attention import ( BlockKVCache, ContextParallelAttention, @@ -34,6 +44,16 @@ from flashdreams.core.attention.rope import apply_rope_freqs +class AttentionBackend(str, Enum): + """Attention implementation used by a Wan DiT block.""" + + WAN = "wan" + """Use Wan's context-parallel attention implementation.""" + + TRITON = "triton" + """Use FP8-projected Triton self- and text cross-attention.""" + + def sinusoidal_embedding_1d(dim: int, position: Tensor) -> Tensor: """Create 1D sinusoidal embeddings. @@ -252,7 +272,7 @@ def update_kv( """Append K/V computed from ``x`` into an existing ``kv_cache``.""" return self._compute_or_update_kv_cache(x, kv_cache, rope_freqs) - def apply_kv( + def query_kv( self, x: Tensor, kv_cache: BlockKVCache, @@ -336,13 +356,13 @@ def forward( rope_freqs_q, rope_freqs_k = self._slice_rope_freqs(rope_freqs, kv_cache) if update_kv_cache: kv_cache = self.update_kv(x, kv_cache, rope_freqs_k) - return self.apply_kv(x, kv_cache, rope_freqs_q, rope_freqs_k) + return self.query_kv(x, kv_cache, rope_freqs_q, rope_freqs_k) class SelfAttention(MultiHeadAttention): """Self-attention that always refreshes K/V cache from current ``x``.""" - def initialize_cache( + def allocate_kv_cache( self, batch_size: int, chunk_size: int, @@ -351,7 +371,7 @@ def initialize_cache( device: torch.device, dtype: torch.dtype, ) -> BlockKVCache: - """Initialize KV cache for streaming self-attention. + """Allocate a KV cache for streaming self-attention. Args: batch_size: Flattened batch size used by attention. @@ -386,6 +406,127 @@ def forward( return super().forward(x, kv_cache, rope_freqs=rope_freqs, update_kv_cache=True) +class TritonSelfAttention(TritonMultiHeadAttention): + """Accelerated self-attention adapted to Wan's checkpoint contract.""" + + @property + def query_projection(self) -> nn.Linear: + """Return Wan's ``q`` module as the logical query projection.""" + return self.q + + @property + def key_projection(self) -> nn.Linear: + """Return Wan's ``k`` module as the logical key projection.""" + return self.k + + @property + def value_projection(self) -> nn.Linear: + """Return Wan's ``v`` module as the logical value projection.""" + return self.v + + @property + def output_projection(self) -> nn.Linear: + """Return Wan's ``o`` module as the logical output projection.""" + return self.o + + @property + def query_norm(self) -> nn.Module: + """Return Wan's ``norm_q`` query normalization module.""" + return self.norm_q + + @property + def key_norm(self) -> nn.Module: + """Return Wan's ``norm_k`` key normalization module.""" + return self.norm_k + + def __init__( + self, + query_dim: int, + context_dim: int | None = None, + n_heads: int = 8, + head_dim: int = 64, + eps: float = 1e-6, + apply_rope_before_kvcache: bool = True, + cp_method: Literal["ring", "ulysses"] = "ring", + sdpa_backend: SDPABackend = SDPABackend.TRITON, + ) -> None: + """Initialize accelerated attention with Wan projection and RoPE policies. + + Args: + query_dim: Feature dimension of input and output tokens. + context_dim: Self-attention context dimension; ``None`` uses + ``query_dim``. + n_heads: Number of attention heads. + head_dim: Per-head feature dimension. + eps: Epsilon used by Q/K RMS normalization. + apply_rope_before_kvcache: Whether keys receive RoPE before cache writes. + Triton requires ``True``. + cp_method: Context-parallel method retained for constructor compatibility. + sdpa_backend: Scaled-dot-product attention implementation. + + Raises: + ValueError: ``context_dim`` differs from ``query_dim`` or cache-relative + RoPE is requested. + """ + del cp_method + context_dim = query_dim if context_dim is None else context_dim + if context_dim != query_dim: + raise ValueError( + "Triton self-attention requires context_dim to equal query_dim; " + f"got {context_dim} and {query_dim}" + ) + if not apply_rope_before_kvcache: + raise ValueError( + "Triton self-attention requires apply_rope_before_kvcache=True" + ) + + super().__init__( + query_dim=query_dim, + n_heads=n_heads, + head_dim=head_dim, + context_dim=context_dim, + attention_type=AttentionType.SELF_ATTENTION, + qkv_fusion_option=QKVFusionOption.FULL, + qk_norm_eps=eps, + qk_norm_scope=QKNormScope.INNER, + rope_interleaved=True, + use_fp8=True, + sdpa_backend=sdpa_backend, + ) + + # Ordinary assignments register only Wan checkpoint names. The logical + # properties above expose these modules to the shared Triton kernels. + self.q = nn.Linear(self.query_dim, self.inner_dim, bias=True) + self.k = nn.Linear(self.context_dim, self.inner_dim, bias=True) + self.v = nn.Linear(self.context_dim, self.inner_dim, bias=True) + self.o = nn.Linear(self.inner_dim, self.query_dim, bias=True) + self.norm_q = nn.RMSNorm(self.inner_dim, eps=self.qk_norm_eps) + self.norm_k = nn.RMSNorm(self.inner_dim, eps=self.qk_norm_eps) + self._initialize_derived_weights() + + def set_context_parallel_group(self, cp_group: ProcessGroup | None) -> None: + """Reject context parallelism unsupported by Triton attention. + + Args: + cp_group: Context-parallel process group; ``None`` is a no-op. + + Raises: + NotImplementedError: ``cp_group`` is not ``None``. + """ + if cp_group is not None: + raise NotImplementedError( + "The Triton attention backend does not support context parallelism" + ) + + def is_context_parallel_enabled(self) -> bool: + """Return whether context parallelism is enabled.""" + return False + + def context_parallel_size(self) -> int: + """Return the singleton context-parallel world size.""" + return 1 + + @dataclass class CrossAttnCache: """Cache container for cross-attention.""" @@ -482,6 +623,127 @@ def forward( return self.o(out) +class TritonCrossAttention(TritonMultiHeadAttention): + """Static text cross-attention adapted to Wan's checkpoint contract.""" + + @property + def query_projection(self) -> nn.Linear: + """Return Wan's ``q`` module as the logical query projection.""" + return self.q + + @property + def key_projection(self) -> nn.Linear: + """Return Wan's ``k`` module as the logical key projection.""" + return self.k + + @property + def value_projection(self) -> nn.Linear: + """Return Wan's ``v`` module as the logical value projection.""" + return self.v + + @property + def output_projection(self) -> nn.Linear: + """Return Wan's ``o`` module as the logical output projection.""" + return self.o + + @property + def query_norm(self) -> nn.Module: + """Return Wan's ``norm_q`` query normalization module.""" + return self.norm_q + + @property + def key_norm(self) -> nn.Module: + """Return Wan's ``norm_k`` key normalization module.""" + return self.norm_k + + def __init__( + self, + query_dim: int, + context_dim: int | None = None, + n_heads: int = 8, + head_dim: int = 64, + eps: float = 1e-6, + cp_method: Literal["ring", "ulysses"] = "ring", + ) -> None: + """Initialize FP8 Triton FA2 cross-attention with Wan module names. + + Args: + query_dim: Feature dimension of query tokens and projected output. + context_dim: Feature dimension of key/value tokens. ``None`` uses + ``query_dim``. + n_heads: Number of attention heads. + head_dim: Per-head feature dimension. + eps: Epsilon used by Q/K RMS normalization. + cp_method: Ignored context-parallel method retained for constructor + compatibility with Wan attention. + """ + del cp_method + super().__init__( + query_dim=query_dim, + context_dim=context_dim, + n_heads=n_heads, + head_dim=head_dim, + attention_type=AttentionType.CROSS_ATTENTION, + qkv_fusion_option=QKVFusionOption.FUSE_KV, + qk_norm_eps=eps, + qk_norm_scope=QKNormScope.INNER, + rope_interleaved=False, + use_fp8=True, + sdpa_backend=SDPABackend.TRITON, + ) + + self.q = nn.Linear(self.query_dim, self.inner_dim, bias=True) + self.k = nn.Linear(self.context_dim, self.inner_dim, bias=True) + self.v = nn.Linear(self.context_dim, self.inner_dim, bias=True) + self.o = nn.Linear(self.inner_dim, self.query_dim, bias=True) + self.norm_q = nn.RMSNorm(self.inner_dim, eps=self.qk_norm_eps) + self.norm_k = nn.RMSNorm(self.inner_dim, eps=self.qk_norm_eps) + self._initialize_derived_weights() + + def initialize_cache( + self, + context_text: Tensor, + context_img: Tensor | None = None, + ) -> CrossAttnCache: + """Project text context into the static cache used by ``forward``. + + Args: + context_text: Text context tensor shaped ``[..., L_text, D]``. + context_img: Unused image context retained for block compatibility. + + Returns: + Cross-attention cache containing projected text K/V. + """ + del context_img + return CrossAttnCache(text=self.compute_kv(context_text)) + + def forward(self, x: Tensor, kv_cache: CrossAttnCache) -> Tensor: + """Attend from ``x`` to the precomputed static text cache.""" + return super().forward(x, kv_cache.text) + + def set_context_parallel_group(self, cp_group: ProcessGroup | None) -> None: + """Reject context parallelism unsupported by Triton attention. + + Args: + cp_group: Context-parallel process group; ``None`` is a no-op. + + Raises: + NotImplementedError: ``cp_group`` is not ``None``. + """ + if cp_group is not None: + raise NotImplementedError( + "The Triton attention backend does not support context parallelism" + ) + + def is_context_parallel_enabled(self) -> bool: + """Return whether context parallelism is enabled.""" + return False + + def context_parallel_size(self) -> int: + """Return the singleton context-parallel world size.""" + return 1 + + @dataclass class BlockCache: """Per-block cache container for self-attention and cross-attention.""" @@ -513,6 +775,8 @@ def __init__( i2v: bool = False, apply_rope_before_kvcache: bool = True, cp_method: Literal["ring", "ulysses"] = "ring", + attention_backend: AttentionBackend = AttentionBackend.WAN, + sdpa_backend: SDPABackend = SDPABackend.TRITON, ) -> None: super().__init__() self.dim = dim @@ -520,30 +784,55 @@ def __init__( self.num_heads = num_heads self.cross_attn_norm = cross_attn_norm self.eps = eps + self.attention_backend = AttentionBackend(attention_backend) + self.sdpa_backend = SDPABackend(sdpa_backend) # Core submodules self.norm1 = nn.LayerNorm(dim, eps=eps, elementwise_affine=False) - self.self_attn = SelfAttention( - query_dim=dim, - n_heads=num_heads, - head_dim=dim // num_heads, - eps=eps, - apply_rope_before_kvcache=apply_rope_before_kvcache, - cp_method=cp_method, - ) + if self.attention_backend is AttentionBackend.WAN: + self.self_attn = SelfAttention( + query_dim=dim, + n_heads=num_heads, + head_dim=dim // num_heads, + eps=eps, + apply_rope_before_kvcache=apply_rope_before_kvcache, + cp_method=cp_method, + ) + else: + self.self_attn = TritonSelfAttention( + query_dim=dim, + n_heads=num_heads, + head_dim=dim // num_heads, + eps=eps, + apply_rope_before_kvcache=apply_rope_before_kvcache, + cp_method=cp_method, + sdpa_backend=self.sdpa_backend, + ) self.norm3 = ( nn.LayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity() ) - self.cross_attn = CrossAttention( - query_dim=dim, - n_heads=num_heads, - head_dim=dim // num_heads, - i2v=i2v, - eps=eps, - cp_method=cp_method, - ) + # ponytail: Wan I2V has independent text/image softmax branches summed + # before one output projection; keep it native until a dual-cache Triton + # adapter can preserve that ordering. + if self.attention_backend is AttentionBackend.TRITON and not i2v: + self.cross_attn = TritonCrossAttention( + query_dim=dim, + n_heads=num_heads, + head_dim=dim // num_heads, + eps=eps, + cp_method=cp_method, + ) + else: + self.cross_attn = CrossAttention( + query_dim=dim, + n_heads=num_heads, + head_dim=dim // num_heads, + i2v=i2v, + eps=eps, + cp_method=cp_method, + ) self.norm2 = nn.LayerNorm(dim, eps=eps, elementwise_affine=False) self.ffn = nn.Sequential( nn.Linear(dim, ffn_dim), @@ -581,7 +870,7 @@ def initialize_cache( dtype = context_text.dtype return BlockCache( - self_attn=self.self_attn.initialize_cache( + self_attn=self.self_attn.allocate_kv_cache( batch_size, chunk_size, window_size, @@ -597,15 +886,19 @@ def set_context_parallel_group(self, cp_group: ProcessGroup | None) -> None: self.self_attn.set_context_parallel_group(cp_group) def update_parameters_after_loading_checkpoint(self) -> None: - """Squeeze the loaded ``[1, 6, D]`` modulation to ``[6, D]``. + """Finalize parameters that depend on loaded checkpoint weights. - Idempotent. Call once after ``load_state_dict`` so the broadcast - in ``forward`` works for any batch shape rather than just the - leading-1 layout the checkpoint was saved in. + Idempotent. Call once after checkpoint tensors are loaded so the broadcast + in ``forward`` works for any batch shape and fused attention weights + reflect the checkpoint. """ if self._parameters_updated_after_loading_checkpoint: return + if isinstance(self.self_attn, TritonSelfAttention): + self.self_attn._refresh_derived_weights() + if isinstance(self.cross_attn, TritonCrossAttention): + self.cross_attn._refresh_derived_weights() self.modulation.data = self.modulation.data.squeeze(0) self._parameters_updated_after_loading_checkpoint = True diff --git a/flashdreams/flashdreams/recipes/wan/transformer/impl/network.py b/flashdreams/flashdreams/recipes/wan/transformer/impl/network.py index 52af5b0fa..9fc4de89e 100644 --- a/flashdreams/flashdreams/recipes/wan/transformer/impl/network.py +++ b/flashdreams/flashdreams/recipes/wan/transformer/impl/network.py @@ -26,12 +26,14 @@ from torch import Tensor from torch.distributed import ProcessGroup +from flashdreams.accelerated.multi_head_attention_triton import SDPABackend from flashdreams.core.distributed.context_parallel import ( cat_outputs_cp, split_inputs_cp, ) from flashdreams.infra.config import InstantiateConfig from flashdreams.recipes.wan.transformer.impl.modules import ( + AttentionBackend, Block, BlockCache, Head, @@ -108,6 +110,10 @@ class WanDiTNetworkConfig(InstantiateConfig): """If True, apply RoPE to keys before storing them in the KV cache.""" cp_method: Literal["ring", "ulysses"] = "ring" """Context-parallel attention method for transformer attention ops.""" + attention_backend: AttentionBackend = AttentionBackend.WAN + """Self- and text cross-attention implementation used by every block.""" + sdpa_backend: SDPABackend = SDPABackend.TRITON + """SDPA implementation used by accelerated self-attention.""" @dataclass @@ -175,6 +181,8 @@ def __init__(self, config: WanDiTNetworkConfig) -> None: self.patch_embedding_type = config.patch_embedding_type self.apply_rope_before_kvcache = config.apply_rope_before_kvcache self.cp_method = config.cp_method + self.attention_backend = AttentionBackend(config.attention_backend) + self.sdpa_backend = SDPABackend(config.sdpa_backend) # Embedding layers in_dim = config.in_dim + 1 if self.concat_padding_mask else config.in_dim @@ -230,6 +238,8 @@ def _build_block(self, layer_idx: int) -> Block: i2v=self.cross_attn_enable_img, apply_rope_before_kvcache=self.apply_rope_before_kvcache, cp_method=self.cp_method, + attention_backend=self.attention_backend, + sdpa_backend=self.sdpa_backend, ) def set_context_parallel_group(self, cp_group: ProcessGroup | None = None) -> None: diff --git a/flashdreams/pyproject.toml b/flashdreams/pyproject.toml index 5a0636e33..843711c64 100644 --- a/flashdreams/pyproject.toml +++ b/flashdreams/pyproject.toml @@ -36,6 +36,7 @@ dependencies = [ # non-subclassable TypeAliasType in NumPy 2.5 (notably on Python 3.13). "numpy>=1.24,<2.5", "nvidia-ml-py>=12.0", + "nvtx>=0.2.15", "psutil>=7.0", "pyyaml>=6.0", "safetensors>=0.4", diff --git a/flashdreams/tests/accelerated/test_multi_head_attention.py b/flashdreams/tests/accelerated/test_multi_head_attention.py new file mode 100644 index 000000000..821ebbed4 --- /dev/null +++ b/flashdreams/tests/accelerated/test_multi_head_attention.py @@ -0,0 +1,242 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""CPU tests for the abstract multi-head attention interface.""" + +from __future__ import annotations + +from typing import cast + +import pytest +from torch import Tensor, nn + +from flashdreams.accelerated.multi_head_attention import ( + AttentionType, + MultiHeadAttention, + QKNormScope, +) + +pytestmark = pytest.mark.ci_cpu + + +class _Attention(MultiHeadAttention[object]): + """Behavior-free concrete attention for base-contract tests. + + The fixture implements the abstract runtime surface without projection + modules so geometry and policy validation stay isolated from any backend. + """ + + @property + def query_projection(self) -> nn.Linear: + """Reject access to the intentionally absent query projection.""" + raise NotImplementedError + + @property + def key_projection(self) -> nn.Linear: + """Reject access to the intentionally absent key projection.""" + raise NotImplementedError + + @property + def value_projection(self) -> nn.Linear: + """Reject access to the intentionally absent value projection.""" + raise NotImplementedError + + @property + def output_projection(self) -> nn.Linear: + """Reject access to the intentionally absent output projection.""" + raise NotImplementedError + + @property + def query_norm(self) -> nn.Module: + """Reject access to the intentionally absent query normalization.""" + raise NotImplementedError + + @property + def key_norm(self) -> nn.Module: + """Reject access to the intentionally absent key normalization.""" + raise NotImplementedError + + def compute_kv( + self, + context: Tensor, + rope_freqs: Tensor | None = None, + ) -> object: + """Return an opaque cache placeholder for interface tests. + + Args: + context: Unused context tokens accepted by the abstract contract. + rope_freqs: Unused positional data; ``None`` is also accepted. + + Returns: + Fresh opaque object standing in for a backend cache. + """ + del context, rope_freqs + return object() + + def forward( + self, + query: Tensor, + kv_cache: object, + rope_freqs: Tensor | None = None, + ) -> Tensor: + """Return ``query`` unchanged through the abstract runtime contract. + + Args: + query: Input tokens passed through unchanged. + kv_cache: Unused opaque cache placeholder. + rope_freqs: Unused positional data; ``None`` is also accepted. + + Returns: + Original ``query`` tensor. + """ + del kv_cache, rope_freqs + return query + + +def test_forward_and_logical_modules_are_abstract() -> None: + """Require one runtime entry point and backend-owned logical modules.""" + # Pin the complete abstract surface so subclasses cannot silently inherit a + # backend-specific projection or cache operation. + assert MultiHeadAttention.__abstractmethods__ == { + "compute_kv", + "forward", + "key_norm", + "key_projection", + "output_projection", + "query_norm", + "query_projection", + "value_projection", + } + + # Cache writes and queries remain private implementation details; callers + # interact with attention only through ``forward``. + assert "query_kv" not in MultiHeadAttention.__dict__ + assert "update_kv" not in MultiHeadAttention.__dict__ + + +def test_attention_geometry_and_policies_are_stored() -> None: + """Store default self-attention and explicit cross-attention policies.""" + # Contrast the default square self-attention geometry with an asymmetric + # cross-attention instance that overrides every optional policy. + self_attention = _Attention( + query_dim=48, + n_heads=3, + head_dim=16, + ) + cross_attention = _Attention( + query_dim=48, + n_heads=3, + head_dim=16, + context_dim=32, + attention_type=AttentionType.CROSS_ATTENTION, + qk_norm_eps=1e-5, + qk_norm_scope=QKNormScope.INNER, + rope_interleaved=True, + ) + + # Self-attention derives context width and inner width from query geometry. + assert self_attention.context_dim == 48 + assert self_attention.inner_dim == 48 + assert self_attention.attention_type is AttentionType.SELF_ATTENTION + assert self_attention.qk_norm_eps == 1e-6 + assert self_attention.qk_norm_scope is QKNormScope.HEAD + assert self_attention.rope_interleaved is False + + # Cross-attention preserves its independent context width and explicit + # normalization and RoPE choices. + assert cross_attention.context_dim == 32 + assert cross_attention.attention_type is AttentionType.CROSS_ATTENTION + assert cross_attention.qk_norm_eps == 1e-5 + assert cross_attention.qk_norm_scope is QKNormScope.INNER + assert cross_attention.rope_interleaved is True + + +def test_attention_type_constructor_policy() -> None: + """Expose stable enum values and reject invalid attention geometry.""" + # Enum values form configuration-facing strings and must remain stable. + assert AttentionType.SELF_ATTENTION.value == "self_attention" + assert AttentionType.CROSS_ATTENTION.value == "cross_attention" + assert QKNormScope.NONE.value == "none" + + # Reject string lookalikes before backend construction can select the wrong + # forward branch. + string_policy = cast(AttentionType, "self_attention") + with pytest.raises(TypeError, match="AttentionType"): + _Attention( + query_dim=16, + n_heads=2, + head_dim=8, + attention_type=string_policy, + ) + + # Self-attention projects one token source, so query and context widths must + # be identical. + with pytest.raises(ValueError, match="self-attention requires"): + _Attention( + query_dim=16, + context_dim=8, + n_heads=2, + head_dim=8, + ) + + +@pytest.mark.parametrize( + "qk_norm_eps", + [ + pytest.param(0.0, id="zero"), + pytest.param(-1.0, id="negative"), + pytest.param(float("inf"), id="infinite"), + pytest.param(float("-inf"), id="negative-infinite"), + pytest.param(float("nan"), id="nan"), + ], +) +def test_qk_norm_eps_must_be_finite_and_positive(qk_norm_eps: float) -> None: + """Reject non-positive or non-finite Q/K normalization epsilon values. + + Args: + qk_norm_eps: Invalid epsilon supplied by the parameterized case. + """ + with pytest.raises(ValueError, match="qk_norm_eps"): + _Attention( + query_dim=16, + n_heads=2, + head_dim=8, + qk_norm_eps=qk_norm_eps, + ) + + +@pytest.mark.parametrize( + ("dimension", "value"), + [ + pytest.param("query_dim", 0, id="zero-query-dim"), + pytest.param("context_dim", 0, id="zero-context-dim"), + pytest.param("n_heads", 0, id="zero-heads"), + pytest.param("head_dim", 0, id="zero-head-dim"), + ], +) +def test_attention_dimensions_must_be_positive(dimension: str, value: int) -> None: + """Reject a zero value for each attention geometry dimension. + + Args: + dimension: Constructor argument under validation. + value: Invalid zero dimension supplied to that argument. + """ + with pytest.raises(ValueError, match=dimension): + _Attention( + query_dim=value if dimension == "query_dim" else 16, + context_dim=value if dimension == "context_dim" else None, + n_heads=value if dimension == "n_heads" else 2, + head_dim=value if dimension == "head_dim" else 8, + ) diff --git a/flashdreams/tests/accelerated/test_multi_head_attention_triton.py b/flashdreams/tests/accelerated/test_multi_head_attention_triton.py new file mode 100644 index 000000000..dbb214a7f --- /dev/null +++ b/flashdreams/tests/accelerated/test_multi_head_attention_triton.py @@ -0,0 +1,367 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Numerical correctness tests for Triton multi-head attention.""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn.functional as F +from torch import Tensor + +from flashdreams.accelerated.multi_head_attention import ( + AttentionType, + QKNormScope, +) +from flashdreams.accelerated.multi_head_attention_triton import ( + QKVFusionOption, + SDPABackend, + TritonMultiHeadAttention, +) + +pytestmark = pytest.mark.ci_gpu + +_QUERY_DIM = 128 +_N_HEADS = 2 +_HEAD_DIM = _QUERY_DIM // _N_HEADS +_CHUNK_SIZE = 16 +_WINDOW_SIZE = 32 +_SINK_SIZE = 4 + + +class _TritonMultiHeadAttention(TritonMultiHeadAttention): + """Checkpoint-compatible Triton attention used by correctness tests.""" + + @property + def query_projection(self) -> torch.nn.Linear: + """Return the query projection.""" + return self.q_proj + + @property + def key_projection(self) -> torch.nn.Linear: + """Return the key projection.""" + return self.k_proj + + @property + def value_projection(self) -> torch.nn.Linear: + """Return the value projection.""" + return self.v_proj + + @property + def output_projection(self) -> torch.nn.Linear: + """Return the output projection.""" + return self.output_proj + + @property + def query_norm(self) -> torch.nn.Module: + """Return the query normalization module.""" + return self.q_norm + + @property + def key_norm(self) -> torch.nn.Module: + """Return the key normalization module.""" + return self.k_norm + + def __init__( + self, + *, + attention_type: AttentionType, + qk_norm_scope: QKNormScope, + rope_interleaved: bool, + sdpa_backend: SDPABackend, + qkv_fusion_option: QKVFusionOption, + use_fp8: bool, + bias: bool, + ) -> None: + """Initialize one correctness-test configuration.""" + super().__init__( + query_dim=_QUERY_DIM, + n_heads=_N_HEADS, + head_dim=_HEAD_DIM, + attention_type=attention_type, + qkv_fusion_option=qkv_fusion_option, + qk_norm_scope=qk_norm_scope, + rope_interleaved=rope_interleaved, + use_fp8=use_fp8, + sdpa_backend=sdpa_backend, + ) + self.q_proj = torch.nn.Linear(_QUERY_DIM, self.inner_dim, bias=bias) + self.k_proj = torch.nn.Linear(_QUERY_DIM, self.inner_dim, bias=bias) + self.v_proj = torch.nn.Linear(_QUERY_DIM, self.inner_dim, bias=bias) + self.output_proj = torch.nn.Linear(self.inner_dim, _QUERY_DIM, bias=bias) + if qk_norm_scope is QKNormScope.NONE: + self.q_norm = torch.nn.Identity() + self.k_norm = torch.nn.Identity() + else: + norm_dim = ( + self.head_dim if qk_norm_scope is QKNormScope.HEAD else self.inner_dim + ) + self.q_norm = torch.nn.RMSNorm(norm_dim, eps=self.qk_norm_eps) + self.k_norm = torch.nn.RMSNorm(norm_dim, eps=self.qk_norm_eps) + self._initialize_derived_weights() + + +@pytest.fixture(scope="module") +def tma_device() -> torch.device: + """Return a CUDA device with tensor-memory acceleration.""" + if not torch.cuda.is_available(): + pytest.skip("CUDA required.") + device = torch.device("cuda") + if torch.cuda.get_device_capability(device)[0] < 9: + pytest.skip("TMA attention requires compute capability 9.0 or newer.") + return device + + +def _normalize( + x: Tensor, + norm: torch.nn.Module, + scope: QKNormScope, +) -> Tensor: + """Apply the configured RMS normalization with native PyTorch math.""" + if scope is QKNormScope.NONE: + return x + assert isinstance(norm, torch.nn.RMSNorm) + original_shape = x.shape + if scope is QKNormScope.INNER: + x = x.flatten(-2) + return F.rms_norm(x, norm.normalized_shape, norm.weight, norm.eps).reshape( + original_shape + ) + + +def _project_query(attention: _TritonMultiHeadAttention, x: Tensor) -> Tensor: + """Project queries with independent PyTorch operators.""" + query = F.linear(x, attention.q_proj.weight, attention.q_proj.bias).reshape( + -1, x.shape[-2], attention.n_heads, attention.head_dim + ) + return _normalize(query, attention.q_norm, attention.qk_norm_scope) + + +def _project_kv( + attention: _TritonMultiHeadAttention, + x: Tensor, +) -> tuple[Tensor, Tensor]: + """Project keys and values with independent PyTorch operators.""" + head_shape = (-1, x.shape[-2], attention.n_heads, attention.head_dim) + key = F.linear(x, attention.k_proj.weight, attention.k_proj.bias).reshape( + head_shape + ) + value = F.linear(x, attention.v_proj.weight, attention.v_proj.bias).reshape( + head_shape + ) + return _normalize(key, attention.k_norm, attention.qk_norm_scope), value + + +def _apply_rope(x: Tensor, rope_freqs: Tensor, interleaved: bool) -> Tensor: + """Apply rotary embeddings with independent PyTorch operators.""" + freqs = rope_freqs[:, 0, 0, :].reshape(1, x.shape[-3], 1, x.shape[-1]) + if interleaved: + rotated = torch.stack((-x[..., 1::2], x[..., 0::2]), dim=-1).flatten(-2) + else: + first, second = x.chunk(2, dim=-1) + rotated = torch.cat((-second, first), dim=-1) + return x * freqs.cos().to(x.dtype) + rotated * freqs.sin().to(x.dtype) + + +def _reference_output( + attention: _TritonMultiHeadAttention, + query: Tensor, + key: Tensor, + value: Tensor, + output_shape: torch.Size, +) -> Tensor: + """Compute non-causal MHA with PyTorch's math SDPA backend.""" + with torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.MATH): + output = F.scaled_dot_product_attention( + query.transpose(1, 2), + key.transpose(1, 2), + value.transpose(1, 2), + dropout_p=0.0, + is_causal=False, + ).transpose(1, 2) + output = F.linear( + output.flatten(-2), + attention.output_proj.weight, + attention.output_proj.bias, + ) + return output.reshape(output_shape) + + +def _rope_freqs( + length: int, + generator: torch.Generator, + device: torch.device, +) -> Tensor: + """Generate deterministic rotary angles for one token sequence.""" + return torch.randn( + length, + 1, + 1, + _HEAD_DIM, + generator=generator, + device=device, + ) + + +def _assert_close(actual: Tensor, expected: Tensor, use_fp8: bool) -> None: + """Compare native or FP8 attention at its numerical error bound.""" + tolerance = 8e-2 if use_fp8 else 2e-2 + torch.testing.assert_close(actual, expected, atol=tolerance, rtol=tolerance) + + +def _check_self_attention( + attention: _TritonMultiHeadAttention, + generator: torch.Generator, + device: torch.device, + dtype: torch.dtype, + use_rope: bool, +) -> None: + """Check streaming self-attention through cache fill and rolling.""" + cache = attention.allocate_kv_cache( + batch_size=1, + chunk_size=_CHUNK_SIZE, + window_size=_WINDOW_SIZE, + sink_size=_SINK_SIZE, + device=device, + dtype=dtype, + ) + all_keys: list[Tensor] = [] + all_values: list[Tensor] = [] + + for chunk_idx in range(3): + x = torch.randn( + 1, + _CHUNK_SIZE, + _QUERY_DIM, + generator=generator, + device=device, + dtype=dtype, + ) + rope = _rope_freqs(_CHUNK_SIZE, generator, device) if use_rope else None + query = _project_query(attention, x) + key, value = _project_kv(attention, x) + if rope is not None: + query = _apply_rope(query, rope, attention.rope_interleaved) + key = _apply_rope(key, rope, attention.rope_interleaved) + all_keys.append(key) + all_values.append(value) + + visible_key = torch.cat(all_keys, dim=1) + visible_value = torch.cat(all_values, dim=1) + if visible_key.shape[1] > _SINK_SIZE + _WINDOW_SIZE: + visible_key = torch.cat( + (visible_key[:, :_SINK_SIZE], visible_key[:, -_WINDOW_SIZE:]), + dim=1, + ) + visible_value = torch.cat( + (visible_value[:, :_SINK_SIZE], visible_value[:, -_WINDOW_SIZE:]), + dim=1, + ) + expected = _reference_output( + attention, query, visible_key, visible_value, x.shape + ) + + cache.before_update(chunk_idx) + actual = attention(x, cache, rope) + cache.after_update(chunk_idx) + _assert_close(actual, expected, attention.use_fp8) + + +def _check_cross_attention( + attention: _TritonMultiHeadAttention, + generator: torch.Generator, + device: torch.device, + dtype: torch.dtype, + use_rope: bool, +) -> None: + """Check static cross-attention against independent PyTorch math.""" + context = torch.randn( + 1, + 24, + _QUERY_DIM, + generator=generator, + device=device, + dtype=dtype, + ) + query_tokens = torch.randn( + 1, + 8, + _QUERY_DIM, + generator=generator, + device=device, + dtype=dtype, + ) + context_rope = _rope_freqs(24, generator, device) if use_rope else None + query_rope = _rope_freqs(8, generator, device) if use_rope else None + + query = _project_query(attention, query_tokens) + key, value = _project_kv(attention, context) + if query_rope is not None and context_rope is not None: + query = _apply_rope(query, query_rope, attention.rope_interleaved) + key = _apply_rope(key, context_rope, attention.rope_interleaved) + expected = _reference_output(attention, query, key, value, query_tokens.shape) + + cache = attention.compute_kv(context, context_rope) + actual = attention(query_tokens, cache, query_rope) + _assert_close(actual, expected, attention.use_fp8) + + +@pytest.mark.parametrize( + "attention_type", tuple(AttentionType), ids=lambda value: value.value +) +@pytest.mark.parametrize( + "qk_norm_scope", tuple(QKNormScope), ids=lambda value: value.value +) +@pytest.mark.parametrize( + "rope_interleaved", [False, True], ids=["rope-split", "rope-interleaved"] +) +@pytest.mark.parametrize( + "sdpa_backend", tuple(SDPABackend), ids=lambda value: value.value +) +@pytest.mark.parametrize( + "qkv_fusion_option", tuple(QKVFusionOption), ids=lambda value: value.value +) +@pytest.mark.parametrize("use_fp8", [False, True], ids=["native", "fp8"]) +@torch.inference_mode() +def test_multi_head_attention_matches_pytorch( + tma_device: torch.device, + attention_type: AttentionType, + qk_norm_scope: QKNormScope, + rope_interleaved: bool, + sdpa_backend: SDPABackend, + qkv_fusion_option: QKVFusionOption, + use_fp8: bool, +) -> None: + """Match every supported MHA policy combination with PyTorch math.""" + dtype = torch.float16 if rope_interleaved else torch.bfloat16 + torch.manual_seed(7) + attention = _TritonMultiHeadAttention( + attention_type=attention_type, + qk_norm_scope=qk_norm_scope, + rope_interleaved=rope_interleaved, + sdpa_backend=sdpa_backend, + qkv_fusion_option=qkv_fusion_option, + use_fp8=use_fp8, + bias=rope_interleaved, + ).to(device=tma_device, dtype=dtype) + attention.eval() + + generator = torch.Generator(device=tma_device).manual_seed(11) + use_rope = qk_norm_scope is not QKNormScope.NONE + if attention_type is AttentionType.SELF_ATTENTION: + _check_self_attention(attention, generator, tma_device, dtype, use_rope) + else: + _check_cross_attention(attention, generator, tma_device, dtype, use_rope) diff --git a/flashdreams/tests/accelerated/triton/conftest.py b/flashdreams/tests/accelerated/triton/conftest.py new file mode 100644 index 000000000..e1e438c23 --- /dev/null +++ b/flashdreams/tests/accelerated/triton/conftest.py @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared fixtures for Triton kernel tests.""" + +from __future__ import annotations + +import pytest +import torch + + +@pytest.fixture(scope="module") +def tma_device() -> torch.device: + """Provide a CUDA device capable of launching TMA kernels. + + Returns: + Active CUDA device with compute capability 9.0 or newer. + """ + # Gate the shared fixture once so every TMA test skips consistently on CPU + # or pre-Hopper hosts instead of failing during kernel compilation. + if not torch.cuda.is_available(): + pytest.skip("CUDA required.") + device = torch.device("cuda") + if torch.cuda.get_device_capability(device)[0] < 9: + pytest.skip("TMA kernels require compute capability 9.0 or newer.") + return device diff --git a/flashdreams/tests/accelerated/triton/test_flash_attention.py b/flashdreams/tests/accelerated/triton/test_flash_attention.py new file mode 100644 index 000000000..583d57a51 --- /dev/null +++ b/flashdreams/tests/accelerated/triton/test_flash_attention.py @@ -0,0 +1,92 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Reference tests for non-causal Triton TMA FlashAttention2 kernels.""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn.functional as F + +from flashdreams.accelerated.triton import flash_attention_2_tma + +pytestmark = pytest.mark.ci_gpu + + +@pytest.mark.parametrize( + ("query_length", "key_length", "head_dim"), + [ + pytest.param(37, 53, 64, id="partial-tiles"), + pytest.param(129, 128, 128, id="production-head-divisible-key"), + ], +) +def test_tma_flash_attention_matches_sdpa( + tma_device: torch.device, + query_length: int, + key_length: int, + head_dim: int, +) -> None: + """Match TMA FlashAttention2 with PyTorch non-causal SDPA. + + Exercise ragged sequence tiles and a production-sized head dimension while + preserving the public token-major ``[B, L, H, D]`` layout. + + Args: + tma_device: CUDA device satisfying the shared TMA capability gate. + query_length: Number of query tokens. + key_length: Number of key and value tokens. + head_dim: Feature width of each attention head. + """ + generator = torch.Generator(device=tma_device).manual_seed(123) + # Generate token-major Q/K/V tensors; Triton consumes and returns this layout. + query = torch.randn( + 1, + query_length, + 2, + head_dim, + generator=generator, + device=tma_device, + dtype=torch.bfloat16, + ) + key = torch.randn( + 1, + key_length, + 2, + head_dim, + generator=generator, + device=tma_device, + dtype=torch.bfloat16, + ) + value = torch.randn( + key.shape, + generator=generator, + device=tma_device, + dtype=torch.bfloat16, + ) + + actual = flash_attention_2_tma(query, key, value) + # PyTorch SDPA consumes head-major ``[B, H, L/S, D]`` views, so transpose + # around the reference call without changing the public comparison layout. + expected = F.scaled_dot_product_attention( + query.transpose(1, 2), + key.transpose(1, 2), + value.transpose(1, 2), + dropout_p=0.0, + is_causal=False, + ).transpose(1, 2) + + # Allow BF16 and online-softmax reduction-order differences between kernels. + torch.testing.assert_close(actual, expected, atol=1e-2, rtol=1e-2) diff --git a/flashdreams/tests/accelerated/triton/test_fp8_quantization.py b/flashdreams/tests/accelerated/triton/test_fp8_quantization.py new file mode 100644 index 000000000..75c010064 --- /dev/null +++ b/flashdreams/tests/accelerated/triton/test_fp8_quantization.py @@ -0,0 +1,81 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Reference tests for Triton row-wise FP8 quantization kernels.""" + +from __future__ import annotations + +import pytest +import torch + +from flashdreams.accelerated.triton import fp8_quantization + +pytestmark = pytest.mark.ci_gpu + + +def test_fused_fp8_row_quantization_matches_torch( + tma_device: torch.device, +) -> None: + """Match fused row-wise E4M3 quantization against PyTorch reference math. + + Exercise a positively strided input, the minimum scale for an all-zero row, + NaN propagation, and an empty row axis that must not launch a kernel. + + Args: + tma_device: CUDA device satisfying the shared TMA capability gate. + """ + generator = torch.Generator(device=tma_device).manual_seed(5) + # Transpose an owned ``[1536, 7]`` allocation to exercise a non-contiguous + # ``[7, 1536]`` view, then reserve row zero for the scale-floor check. + x = torch.randn( + (1536, 7), + generator=generator, + device=tma_device, + dtype=torch.bfloat16, + ).T + x[0].zero_() + actual, actual_scales = fp8_quantization._quantize_fp8_rows(x) + + # Reproduce the kernel's per-row scale and normalized E4M3 cast in FP32. + x_float = x.to(torch.float32) + expected_scales = ( + (x_float.abs().amax(dim=1, keepdim=True) / fp8_quantization._FP8_MAX) + .clamp_min(1e-12) + .contiguous() + ) + expected = ( + (x_float / expected_scales) + .clamp(-fp8_quantization._FP8_MAX, fp8_quantization._FP8_MAX) + .to(torch.float8_e4m3fn) + ) + + # Exact arithmetic parity is expected, and contiguous output is part of the + # wrapper contract even when the source view is strided. + assert actual.is_contiguous() + assert torch.equal(actual, expected) + assert torch.equal(actual_scales, expected_scales) + assert torch.count_nonzero(actual[0]) == 0 + assert actual_scales[0, 0] == 1e-12 + + # Pin IEEE NaN propagation separately from finite-row equality. + nan_input = x.clone() + nan_input[0, 0] = float("nan") + nan_output, _ = fp8_quantization._quantize_fp8_rows(nan_input) + assert torch.isnan(nan_output[0, 0].to(torch.float32)) + + # An empty row axis returns correctly shaped storage without a Triton launch. + empty_output, empty_scales = fp8_quantization._quantize_fp8_rows(x[:0]) + assert empty_output.shape == x[:0].shape + assert empty_scales.shape == (0, 1) diff --git a/flashdreams/tests/accelerated/triton/test_rms_rope_kv_cache.py b/flashdreams/tests/accelerated/triton/test_rms_rope_kv_cache.py new file mode 100644 index 000000000..f25e90605 --- /dev/null +++ b/flashdreams/tests/accelerated/triton/test_rms_rope_kv_cache.py @@ -0,0 +1,191 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Reference tests for fused Triton RMSNorm, RoPE, and K/V-cache updates.""" + +from __future__ import annotations + +import pytest +import torch + +from flashdreams.accelerated.multi_head_attention import QKNormScope +from flashdreams.accelerated.triton import fused_rms_rope_kv_cache_update + +pytestmark = pytest.mark.ci_gpu + + +def _apply_rope_reference( + x: torch.Tensor, rope_freqs: torch.Tensor, *, interleaved: bool +) -> torch.Tensor: + """Apply full-width RoPE with independent PyTorch tensor operations. + + Compute ``x * cos(theta) + rotate(x) * sin(theta)`` using either adjacent + feature pairs or matching positions from the two head-dimension halves. + + Args: + x: Head activations with shape ``[B, L, H, D]``. + rope_freqs: Rotation angles with shape ``[L, 1, 1, D]``. + interleaved: Pair adjacent features when ``True``; otherwise pair the + first and second halves of each head. + + Returns: + Rotated activations with the same shape and dtype as ``x``. + """ + # Reshape ``[L, 1, 1, D]`` angles to broadcast across batch and head axes. + freqs = rope_freqs[:, 0, 0, :].reshape(1, x.shape[-3], 1, x.shape[-1]) + cos_freqs = freqs.cos().to(dtype=x.dtype) + sin_freqs = freqs.sin().to(dtype=x.dtype) + if interleaved: + rotated = torch.stack((-x[..., 1::2], x[..., 0::2]), dim=-1).flatten(-2) + else: + first, second = x.chunk(2, dim=-1) + rotated = torch.cat((-second, first), dim=-1) + return x * cos_freqs + rotated * sin_freqs + + +@pytest.mark.parametrize( + ("qk_norm_scope", "rope_interleaved", "head_major_cache", "apply_rope"), + [ + pytest.param(QKNormScope.HEAD, False, True, True, id="half-split-head-major"), + pytest.param(QKNormScope.HEAD, True, True, True, id="interleaved-head-major"), + pytest.param(QKNormScope.NONE, False, False, True, id="no-norm-token-major"), + pytest.param(QKNormScope.HEAD, False, True, False, id="no-rope-head-major"), + ], +) +def test_head_preprocessing_layouts_match_torch( + tma_device: torch.device, + qk_norm_scope: QKNormScope, + rope_interleaved: bool, + head_major_cache: bool, + apply_rope: bool, +) -> None: + """Match fused head preprocessing and sliced cache writes against PyTorch. + + Cover head-scoped and disabled normalization, both RoPE pair conventions, + optional RoPE, and token-major or physically head-major cache storage. + + Args: + tma_device: CUDA device satisfying the shared TMA capability gate. + qk_norm_scope: RMS normalization scope exercised by the kernel. + rope_interleaved: RoPE feature-pair convention. + head_major_cache: Store cache data physically as ``[B, H, S, D]`` while + exposing the logical ``[B, S, H, D]`` interface. + apply_rope: Apply rotation when ``True``. + """ + batch_size, sequence_length, num_heads, head_dim = 2, 3, 3, 16 + cache_size = 6 + generator = torch.Generator(device=tma_device).manual_seed(23) + query = torch.randn( + batch_size, + sequence_length, + num_heads, + head_dim, + generator=generator, + device=tma_device, + dtype=torch.bfloat16, + ) + key = torch.randn( + query.shape, + generator=generator, + device=tma_device, + dtype=torch.bfloat16, + ) + value = torch.randn( + query.shape, + generator=generator, + device=tma_device, + dtype=torch.bfloat16, + ) + sentinel = -123.0 + # Build logical ``[B, S, H, D]`` caches in either contiguous token-major + # storage or a transposed view over contiguous ``[B, H, S, D]`` storage. + cache_shape = ( + (batch_size, num_heads, cache_size, head_dim) + if head_major_cache + else (batch_size, cache_size, num_heads, head_dim) + ) + key_cache = torch.full( + cache_shape, + sentinel, + device=tma_device, + dtype=torch.bfloat16, + ) + if head_major_cache: + key_cache = key_cache.transpose(1, 2) + value_cache = torch.full_like(key_cache, sentinel) + rope_freqs = ( + torch.arange( + sequence_length * head_dim, + device=tma_device, + dtype=torch.float32, + ).reshape(sequence_length, 1, 1, head_dim) + / 37 + ) + query_weight = key_weight = None + if qk_norm_scope is QKNormScope.HEAD: + query_weight = torch.ones(head_dim, device=tma_device, dtype=torch.bfloat16) + key_weight = torch.ones_like(query_weight) + + # Process every query token while copying source tokens ``[1:3]`` into cache + # positions ``[2:4]``; neighboring sentinel rows must remain untouched. + actual_query = fused_rms_rope_kv_cache_update( + query, + key, + value, + key_cache, + value_cache, + query_weight=query_weight, + key_weight=key_weight, + norm_eps=1e-6, + norm_scope=qk_norm_scope, + rope_freqs=rope_freqs if apply_rope else None, + rope_interleaved=rope_interleaved, + cache_read_start=1, + cache_write_start=2, + cache_write_length=2, + ) + + # Compose the reference in kernel order: RMSNorm, then RoPE. Values bypass + # both operations and are copied directly into the selected cache slice. + expected_query = query + expected_key = key + if qk_norm_scope is QKNormScope.HEAD: + expected_query = torch.nn.functional.rms_norm( + query, (head_dim,), weight=query_weight, eps=1e-6 + ) + expected_key = torch.nn.functional.rms_norm( + key, (head_dim,), weight=key_weight, eps=1e-6 + ) + if apply_rope: + expected_query = _apply_rope_reference( + expected_query, rope_freqs, interleaved=rope_interleaved + ) + expected_key = _apply_rope_reference( + expected_key, rope_freqs, interleaved=rope_interleaved + ) + + # Allow BF16 reduction and trigonometric ordering differences for processed + # Q/K; the unmodified V slice must remain exact. + torch.testing.assert_close(actual_query, expected_query, atol=1e-2, rtol=1e-2) + torch.testing.assert_close( + key_cache[:, 2:4], expected_key[:, 1:3], atol=1e-2, rtol=1e-2 + ) + torch.testing.assert_close(value_cache[:, 2:4], value[:, 1:3]) + + # Sentinel integrity proves that the fused store honors both write bounds. + assert torch.all(key_cache[:, :2] == sentinel) + assert torch.all(key_cache[:, 4:] == sentinel) + assert torch.all(value_cache[:, :2] == sentinel) + assert torch.all(value_cache[:, 4:] == sentinel) diff --git a/flashdreams/tests/recipes/test_wan_modules.py b/flashdreams/tests/recipes/test_wan_modules.py new file mode 100644 index 000000000..a1bc80c4d --- /dev/null +++ b/flashdreams/tests/recipes/test_wan_modules.py @@ -0,0 +1,218 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""GPU correctness tests for accelerated Wan recipe components.""" + +from __future__ import annotations + +import pytest +import torch + +from flashdreams.accelerated.multi_head_attention_triton import ( + SDPABackend, + TritonMultiHeadAttention, +) +from flashdreams.core.attention import RotaryPositionEmbedding3D +from flashdreams.recipes.wan.transformer.impl.modules import ( + AttentionBackend, + Block, + CrossAttention, + SelfAttention, + TritonCrossAttention, +) + +pytestmark = pytest.mark.ci_gpu + + +@pytest.fixture(scope="module") +def tma_device() -> torch.device: + """Return a CUDA device that supports tensor-memory acceleration.""" + if not torch.cuda.is_available(): + pytest.skip("CUDA required.") + device = torch.device("cuda") + if torch.cuda.get_device_capability(device)[0] < 9: + pytest.skip("TMA attention requires compute capability 9.0 or newer.") + return device + + +@pytest.mark.parametrize( + "sdpa_backend", tuple(SDPABackend), ids=lambda backend: backend.value +) +def test_triton_self_attention_matches_default_through_window_roll( + tma_device: torch.device, + sdpa_backend: SDPABackend, +) -> None: + """Compare Triton with Wan self-attention through cache fill and roll.""" + torch.manual_seed(7) + default_block = Block( + dim=256, + ffn_dim=512, + num_heads=2, + ) + triton_block = Block( + dim=256, + ffn_dim=512, + num_heads=2, + attention_backend=AttentionBackend.TRITON, + sdpa_backend=sdpa_backend, + ) + triton_block.load_state_dict(default_block.state_dict(), strict=True) + + assert default_block.attention_backend is AttentionBackend.WAN + reference = default_block.self_attn.to( + device=tma_device, dtype=torch.bfloat16 + ).eval() + actual_attention = triton_block.self_attn.to( + device=tma_device, dtype=torch.bfloat16 + ).eval() + assert isinstance(reference, SelfAttention) + assert isinstance(actual_attention, TritonMultiHeadAttention) + assert actual_attention.use_fp8 is True + assert actual_attention.sdpa_backend is sdpa_backend + + batch_size = 1 + len_t, len_h, len_w = 1, 1, 16 + chunk_size = len_t * len_h * len_w + window_size = 2 * chunk_size + reference_cache = reference.allocate_kv_cache( + batch_size=batch_size, + chunk_size=chunk_size, + window_size=window_size, + sink_size=0, + device=tma_device, + dtype=torch.bfloat16, + ) + actual_cache = actual_attention.allocate_kv_cache( + batch_size=batch_size, + chunk_size=chunk_size, + window_size=window_size, + sink_size=0, + device=tma_device, + dtype=torch.bfloat16, + ) + assert actual_cache.dtype is ( + torch.float8_e4m3fn if sdpa_backend is SDPABackend.TRITON else torch.bfloat16 + ) + rope = RotaryPositionEmbedding3D( + head_dim=128, + len_t=len_t, + len_h=len_h, + len_w=len_w, + interleaved=True, + device=tma_device, + ) + generator = torch.Generator(device=tma_device).manual_seed(11) + + with torch.inference_mode(): + for chunk_idx in range(3): + x = torch.randn( + batch_size, + chunk_size, + 256, + generator=generator, + device=tma_device, + dtype=torch.bfloat16, + ) + rope_freqs = rope.shift_t(chunk_idx) + reference_cache.before_update(chunk_idx) + actual_cache.before_update(chunk_idx) + + expected = reference(x, reference_cache, rope_freqs) + actual = actual_attention(x, actual_cache, rope_freqs) + + torch.testing.assert_close(actual, expected, atol=5e-2, rtol=5e-2) + torch.testing.assert_close( + actual_cache.cached_k().to(torch.bfloat16), + reference_cache.cached_k(), + atol=1.5e-1, + rtol=1.25e-1, + ) + torch.testing.assert_close( + actual_cache.cached_v().to(torch.bfloat16), + reference_cache.cached_v(), + atol=1.5e-1, + rtol=1.25e-1, + ) + + reference_cache.after_update(chunk_idx) + actual_cache.after_update(chunk_idx) + + +def test_triton_cross_attention_matches_default_static_text_cache( + tma_device: torch.device, +) -> None: + """Match Wan T2V cross-attention while preserving static text caches. + + Args: + tma_device: CUDA device with tensor-memory acceleration support. + """ + torch.manual_seed(19) + reference = ( + CrossAttention( + query_dim=256, + n_heads=2, + head_dim=128, + ) + .to( + device=tma_device, + dtype=torch.bfloat16, + ) + .eval() + ) + actual_attention = ( + TritonCrossAttention( + query_dim=256, + n_heads=2, + head_dim=128, + ) + .to( + device=tma_device, + dtype=torch.bfloat16, + ) + .eval() + ) + actual_attention.load_state_dict(reference.state_dict(), strict=True) + + generator = torch.Generator(device=tma_device).manual_seed(23) + context = torch.randn( + (1, 32, 256), + generator=generator, + device=tma_device, + dtype=torch.bfloat16, + ) + query = torch.randn( + (1, 16, 256), + generator=generator, + device=tma_device, + dtype=torch.bfloat16, + ) + + with torch.inference_mode(): + reference_cache = reference.initialize_cache(context) + actual_cache = actual_attention.initialize_cache(context) + reference_key = reference_cache.text.cached_k().clone() + reference_value = reference_cache.text.cached_v().clone() + actual_key = actual_cache.text.cached_k().clone() + actual_value = actual_cache.text.cached_v().clone() + + expected = reference(query, reference_cache) + actual = actual_attention(query, actual_cache) + + torch.testing.assert_close(actual, expected, atol=5e-2, rtol=5e-2) + assert torch.equal(reference_cache.text.cached_k(), reference_key) + assert torch.equal(reference_cache.text.cached_v(), reference_value) + assert torch.equal(actual_cache.text.cached_k(), actual_key) + assert torch.equal(actual_cache.text.cached_v(), actual_value) + assert actual_cache.text.dtype is torch.float8_e4m3fn diff --git a/flashdreams/tests/test_fp8_kvcache.py b/flashdreams/tests/test_fp8_kvcache.py new file mode 100644 index 000000000..303f3b744 --- /dev/null +++ b/flashdreams/tests/test_fp8_kvcache.py @@ -0,0 +1,88 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""CPU lifecycle tests for FP8 block K/V storage.""" + +from __future__ import annotations + +import pytest +import torch + +from flashdreams.core.attention import BlockKVCache + +pytestmark = pytest.mark.ci_cpu + + +def test_fp8_block_kv_cache_converts_native_inputs() -> None: + """Store native-precision inputs in fixed E4M3 cache tensors.""" + key = torch.tensor([[[[1.0, -2.0]], [[3.0, -4.0]]]]) + value = torch.tensor([[[[5.0, -6.0]], [[7.0, -8.0]]]]) + + cache = BlockKVCache( + k_shape=tuple(key.shape), + v_shape=tuple(value.shape), + seq_dim=1, + chunk_size=key.shape[1], + window_size=key.shape[1], + device="cpu", + dtype=torch.float8_e4m3fn, + ) + cache.before_update(0) + cache.update(key, value) + + assert cache._k.dtype is torch.float8_e4m3fn + assert cache._v.dtype is torch.float8_e4m3fn + torch.testing.assert_close(cache.cached_k().float(), key) + torch.testing.assert_close(cache.cached_v().float(), value) + cache.after_update(0) + + +def test_fp8_block_kv_cache_preserves_rolling_lifecycle() -> None: + """Preserve filling, rolling, overwrite, and reset behavior in FP8.""" + cache = BlockKVCache( + k_shape=(1, 4, 1, 1), + v_shape=(1, 4, 1, 1), + seq_dim=1, + chunk_size=2, + window_size=4, + device="cpu", + dtype=torch.float8_e4m3fn, + ) + + for chunk_idx in range(3): + values = torch.tensor([[[[2.0 * chunk_idx]], [[2.0 * chunk_idx + 1.0]]]]) + cache.before_update(chunk_idx) + cache.update(values, -values) + cache.after_update(chunk_idx) + + cache.before_update(2) + replacement = torch.tensor([[[[8.0]], [[9.0]]]]) + cache.update(replacement, -replacement) + + torch.testing.assert_close( + cache.cached_k().float(), + torch.tensor([[[[2.0]], [[3.0]], [[8.0]], [[9.0]]]]), + ) + torch.testing.assert_close( + cache.cached_v().float(), + torch.tensor([[[[-2.0]], [[-3.0]], [[-8.0]], [[-9.0]]]]), + ) + cache.after_update(2) + + key_pointer = cache._k.data_ptr() + cache.reset() + + assert cache._k.data_ptr() == key_pointer + assert cache.size == 0 diff --git a/flashdreams/tests/test_wan_context_parallel.py b/flashdreams/tests/test_wan_context_parallel.py index df772aade..7eff433ed 100644 --- a/flashdreams/tests/test_wan_context_parallel.py +++ b/flashdreams/tests/test_wan_context_parallel.py @@ -8,6 +8,12 @@ import pytest import torch +from flashdreams.accelerated.multi_head_attention import AttentionType +from flashdreams.accelerated.multi_head_attention_triton import ( + QKVFusionOption, + SDPABackend, + TritonMultiHeadAttention, +) from flashdreams.core.attention.kvcache import BlockKVCache from flashdreams.recipes.wan.transformer.impl import modules as wan_modules from flashdreams.recipes.wan.transformer.impl.network import WanDiTNetworkConfig @@ -114,7 +120,7 @@ def _fake_apply_rope_freqs(x, freqs, interleaved=False): return x.add_(1.0) monkeypatch.setattr(wan_modules, "apply_rope_freqs", _fake_apply_rope_freqs) - attn.apply_kv( + attn.query_kv( torch.randn(1, 3, 4), cache, rope_freqs_q=torch.zeros(3, 1, 1, 4), @@ -163,6 +169,100 @@ def test_wan21_requires_tokens_divisible_by_cp_size(monkeypatch) -> None: ) +@pytest.mark.parametrize( + "sdpa_backend", tuple(SDPABackend), ids=lambda backend: backend.value +) +def test_wan_network_propagates_sdpa_backend( + sdpa_backend: SDPABackend, +) -> None: + """Propagate the configured SDPA implementation through every DiT block.""" + config = WanDiTNetworkConfig( + dim=64, + ffn_dim=128, + num_heads=4, + num_layers=1, + patch_embedding_type="linear", + attention_backend=wan_modules.AttentionBackend.TRITON, + sdpa_backend=sdpa_backend, + ) + + network = config.setup() + block = network.blocks[0] + + assert config.sdpa_backend is sdpa_backend + assert network.sdpa_backend is sdpa_backend + assert block.sdpa_backend is sdpa_backend + assert isinstance(block.self_attn, TritonMultiHeadAttention) + assert block.self_attn.qkv_fusion_option is QKVFusionOption.FULL + assert block.self_attn.sdpa_backend is sdpa_backend + assert isinstance(block.cross_attn, wan_modules.TritonCrossAttention) + assert block.cross_attn.attention_type is AttentionType.CROSS_ATTENTION + assert block.cross_attn.qkv_fusion_option is QKVFusionOption.FUSE_KV + assert block.cross_attn.use_fp8 is True + assert block.cross_attn.sdpa_backend is SDPABackend.TRITON + + +def test_triton_attention_registers_only_wan_checkpoint_names() -> None: + """Expose logical Triton modules without adding generic checkpoint aliases.""" + reference = wan_modules.SelfAttention(64, n_heads=4, head_dim=16) + actual = wan_modules.TritonSelfAttention(64, n_heads=4, head_dim=16) + actual.load_state_dict(reference.state_dict(), strict=True) + + assert actual.query_projection is actual.q + assert actual.key_projection is actual.k + assert actual.value_projection is actual.v + assert actual.output_projection is actual.o + assert actual.query_norm is actual.norm_q + assert actual.key_norm is actual.norm_k + + generic_names = {"q_proj", "k_proj", "v_proj", "output_proj", "q_norm", "k_norm"} + assert generic_names.isdisjoint(actual._modules) + assert set(actual.state_dict()) == set(reference.state_dict()) + assert actual._derived_weights.fused_qkv_weight is not None + + +def test_triton_cross_attention_registers_only_wan_checkpoint_names() -> None: + """Load Wan T2V cross-attention weights without generic checkpoint aliases.""" + reference = wan_modules.CrossAttention(query_dim=64, n_heads=4, head_dim=16) + actual = wan_modules.TritonCrossAttention(query_dim=64, n_heads=4, head_dim=16) + actual.load_state_dict(reference.state_dict(), strict=True) + + assert actual.query_projection is actual.q + assert actual.key_projection is actual.k + assert actual.value_projection is actual.v + assert actual.output_projection is actual.o + assert actual.query_norm is actual.norm_q + assert actual.key_norm is actual.norm_k + + generic_names = {"q_proj", "k_proj", "v_proj", "output_proj", "q_norm", "k_norm"} + assert generic_names.isdisjoint(actual._modules) + assert set(actual.state_dict()) == set(reference.state_dict()) + assert actual._derived_weights.fused_kv_weight is not None + + +def test_triton_backend_keeps_native_i2v_cross_attention() -> None: + """Keep Wan's dual text/image cross-attention when I2V is enabled.""" + reference = wan_modules.Block( + dim=64, + ffn_dim=128, + num_heads=4, + i2v=True, + ) + actual = wan_modules.Block( + dim=64, + ffn_dim=128, + num_heads=4, + i2v=True, + attention_backend=wan_modules.AttentionBackend.TRITON, + ) + actual.load_state_dict(reference.state_dict(), strict=True) + + assert isinstance(actual.self_attn, wan_modules.TritonSelfAttention) + assert isinstance(actual.cross_attn, wan_modules.CrossAttention) + assert not isinstance(actual.cross_attn, wan_modules.TritonCrossAttention) + assert actual.cross_attn.i2v is True + + def test_wan_patchify_unpatchify_round_trip_without_cp() -> None: network = WanDiTNetworkConfig( dim=64, diff --git a/integrations/lingbot/README.md b/integrations/lingbot/README.md index e993a539d..f1cf35577 100644 --- a/integrations/lingbot/README.md +++ b/integrations/lingbot/README.md @@ -309,6 +309,13 @@ Text-driven events work with both v1 and v2 through the same DataChannel: `event_catalog`, and `active_event_id`; successful updates receive an `event_ack`. +## Benchmarks + +Manual GPU benchmarks cover the LingBot-owned `CamCtrlBlock`, the complete DiT +network, and steady-state full-pipeline `generate` and `finalize` stages. See +the [benchmark guide](benchmarks/README.md) for hardware requirements, scope, +and the single- and multi-GPU commands. + ## Tests ```bash diff --git a/integrations/lingbot/benchmarks/README.md b/integrations/lingbot/benchmarks/README.md new file mode 100644 index 000000000..6703980b2 --- /dev/null +++ b/integrations/lingbot/benchmarks/README.md @@ -0,0 +1,111 @@ + + +# LingBot benchmarks + +The LingBot benchmarks are manual, GPU-only pytest tests. The complete suite can +run on one sufficiently large GPU or with four-way context parallelism. The +single-GPU path is validated on a 256 GB NVIDIA GB300; smaller devices may run +out of memory during compile/autotune or cache setup. In the distributed path, +the 14B network has about 34.5 GiB of BF16 weights per rank before its KV cache, +activations, and compiler workspaces. The full-pipeline cases also require the +Hugging Face access and cache space documented in +[`integrations/lingbot/README.md`](../README.md). + +Between benchmark cases, pytest synchronizes CUDA, resets compiler and CUDA-graph +state, collects Python objects, and empties the CUDA allocator cache. This prevents +one case from retaining memory needed by the next, but does not reduce a case's own +peak memory requirement. + +Each benchmark layer runs three cases: the WAN/cuDNN reference, Triton FP8 +projections with PyTorch cuDNN SDPA, and Triton FP8 projections with Triton +FlashAttention2 (FA2). Their stable labels are `wan_torch`, `triton_cudnn`, and +`triton_fa2`. Both Triton cases require compute capability 9.0 or newer and run +only on a single GPU because Triton attention does not support context +parallelism; cross-attention remains cuDNN in all three cases. + +First sync the LingBot package and the workspace `test` dependency group, +which provides both `pytest` and `pytest-benchmark`: + +```bash +uv sync --package flashdreams-lingbot --group test +``` + +Run the complete suite on one high-memory GPU: + +```bash +uv run --package flashdreams-lingbot --group test pytest \ + integrations/lingbot/benchmarks \ + -p no:manual_marker -m manual --benchmark-only -v +``` + +Alternatively, run with four-way context parallelism when four usable CUDA +ordinals are visible in the same process namespace: + +```bash +uv run --package flashdreams-lingbot --group test \ + torchrun --standalone --nproc_per_node=4 --no-python pytest \ + integrations/lingbot/benchmarks \ + -p no:manual_marker -m manual --benchmark-only -v +``` + +All four workers execute every WAN test because the DiT uses context-parallel +collectives. The tests align workers before every sample. Triton cases skip in +this distributed mode because that backend is single-GPU only. Once the command +is known to work, add `--local-ranks-filter=0` before `--no-python` to show +only rank 0's report; omit it while debugging because it also hides tracebacks +from failing nonzero ranks. If `torchrun` reports `invalid device ordinal`, +reduce `--nproc_per_node` or scope `CUDA_VISIBLE_DEVICES` to GPUs that are +usable together. Do not point multiple workers at one shared +`--benchmark-json` path; use a rank-specific path when retaining every rank's +raw report. + +The block microbenchmark also fits on one large GPU: + +```bash +uv run --package flashdreams-lingbot --group test pytest \ + integrations/lingbot/benchmarks/test_modules.py \ + -p no:manual_marker -m manual --benchmark-only -v +``` + +To select another benchmark layer, use one of these files: + +- `test_modules.py` benchmarks only LingBot's integration-owned + `CamCtrlBlock`. It deliberately does not add standalone benchmarks for the + inherited Wan attention, MLP, normalization, encoder, or decoder modules. + The whole-block timing necessarily includes the inherited Wan branches that + the LingBot subclass executes. +- `test_network.py` benchmarks one steady-state evaluation of the complete + random-initialized LingBot 14B camera-control DiT. It uses the CLI replay's + 352x640 geometry, compiled WAN or Triton self-attention, CUDA graph replay, + and the shipped window15/sink3 cache layout; startup is excluded. +- `test_pipeline.py` separately benchmarks steady-state `generate` and + `finalize` for + `lingbot-world-v2-14b-causal-fast-taehv-window15-sink3` at the CLI replay's + 352x640 geometry. This is an end-to-end measurement, so the recurring path + includes its configured recipe components, but metadata identifies those + stages and no reused component is presented as a LingBot module + microbenchmark. At the measured AR indices, the Wan I2V VAE branch reuses + its cached latent. Reported output FPS is generate-only; use the separate + finalize result when evaluating the complete per-chunk lifecycle. + +Keep `--group test` on both setup and run commands. A package-only environment +does not include the benchmark plugin. The benchmarks exclude setup, cache +fill, compile/autotune, and CUDA graph capture from measured rounds. When +publishing results, also record the exact command and commit, GPU/driver and +software stack, checkpoint identifiers, compiler-cache state, and fallback +warnings. diff --git a/integrations/lingbot/benchmarks/cases.py b/integrations/lingbot/benchmarks/cases.py new file mode 100644 index 000000000..89daef07a --- /dev/null +++ b/integrations/lingbot/benchmarks/cases.py @@ -0,0 +1,106 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared Lingbot attention benchmark cases.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import pytest +import torch + +from flashdreams.accelerated.multi_head_attention_triton import SDPABackend +from flashdreams.recipes.wan.transformer.impl.modules import AttentionBackend + + +@dataclass(frozen=True) +class AttentionBenchmarkCase: + """Configuration and metadata for one attention benchmark implementation.""" + + implementation: str + """Stable implementation name stored in benchmark metadata.""" + + attention_backend: AttentionBackend + """DiT block implementation configured for this case.""" + + sdpa_backend: SDPABackend + """SDPA implementation configured for Triton self-attention.""" + + self_attention_operator: str + """Self-attention operator reported in benchmark metadata.""" + + cross_attention_operator: str + """Cross-attention operator reported in benchmark metadata.""" + + minimum_compute_capability: tuple[int, int] | None = None + """Minimum CUDA compute capability; ``None`` accepts any CUDA device.""" + + @property + def pytest_id(self) -> str: + """Return the readable pytest parameter identifier.""" + return self.implementation.replace("_", "-") + + +WAN_TORCH_CASE = AttentionBenchmarkCase( + implementation="wan_torch", + attention_backend=AttentionBackend.WAN, + sdpa_backend=SDPABackend.CUDNN, + self_attention_operator="cudnn", + cross_attention_operator="cudnn", +) + +TRITON_CUDNN_CASE = AttentionBenchmarkCase( + implementation="triton_cudnn", + attention_backend=AttentionBackend.TRITON, + sdpa_backend=SDPABackend.CUDNN, + self_attention_operator="torch_cudnn_sdpa", + cross_attention_operator="triton_fa2", + minimum_compute_capability=(9, 0), +) + +TRITON_FA2_CASE = AttentionBenchmarkCase( + implementation="triton_fa2", + attention_backend=AttentionBackend.TRITON, + sdpa_backend=SDPABackend.TRITON, + self_attention_operator="triton_fa2", + cross_attention_operator="triton_fa2", + minimum_compute_capability=(9, 0), +) + +ATTENTION_CASES = (WAN_TORCH_CASE, TRITON_CUDNN_CASE, TRITON_FA2_CASE) +"""Attention cases exercised by each Lingbot benchmark layer.""" + +assert {case.attention_backend for case in ATTENTION_CASES} == set(AttentionBackend) +assert { + case.sdpa_backend + for case in ATTENTION_CASES + if case.attention_backend is AttentionBackend.TRITON +} == set(SDPABackend) + + +def skip_unsupported_device( + case: AttentionBenchmarkCase, + device: torch.device, +) -> None: + """Skip a benchmark case when device is older than its minimum capability.""" + minimum = case.minimum_compute_capability + if minimum is None: + return + if torch.cuda.get_device_capability(device) < minimum: + pytest.skip( + f"{case.pytest_id} attention requires compute capability " + f"{minimum[0]}.{minimum[1]}+" + ) diff --git a/integrations/lingbot/benchmarks/conftest.py b/integrations/lingbot/benchmarks/conftest.py new file mode 100644 index 000000000..753fcab94 --- /dev/null +++ b/integrations/lingbot/benchmarks/conftest.py @@ -0,0 +1,34 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared pytest fixtures for LingBot benchmarks.""" + +import gc +from collections.abc import Iterator + +import pytest +import torch + + +@pytest.fixture(autouse=True) +def _release_cuda_memory_between_benchmarks() -> Iterator[None]: + """Release compiler and CUDA allocator state after each benchmark.""" + yield + if not torch.cuda.is_available(): + return + torch.cuda.synchronize() + torch.compiler.reset() + gc.collect() + torch.cuda.empty_cache() diff --git a/integrations/lingbot/benchmarks/test_modules.py b/integrations/lingbot/benchmarks/test_modules.py new file mode 100644 index 000000000..24ec98b5a --- /dev/null +++ b/integrations/lingbot/benchmarks/test_modules.py @@ -0,0 +1,340 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Microbenchmark for the LingBot-owned camera-control DiT block. + +Run all attention cases with ``uv run --package flashdreams-lingbot +--group test pytest integrations/lingbot/benchmarks/test_modules.py +-p no:manual_marker -m manual --benchmark-only``. +""" + +from __future__ import annotations + +import os + +import pytest +import torch +import torch.distributed as dist +from lingbot.transformer.impl.modules import CamCtrlBlock +from lingbot.transformer.impl.network import LingbotWorldDiTNetwork14BConfig +from pytest_benchmark.fixture import BenchmarkFixture + +from flashdreams.core.attention.rope import RotaryPositionEmbedding3D +from flashdreams.core.distributed import init as init_distributed +from flashdreams.recipes.wan.transformer.impl.modules import AttentionBackend +from integrations.lingbot.benchmarks.cases import ( + ATTENTION_CASES, + AttentionBenchmarkCase, + skip_unsupported_device, +) + +pytestmark = pytest.mark.manual + +_GPU_REASON = "LingBot DiT block benchmark requires CUDA" + +# CLI replay geometry: 352x640 pixels become 44x80 Wan +# latents. The DiT consumes three latent frames per chunk; its window15/sink3 +# preset retains six chunks in total while bounding cache memory. +_PIXEL_HEIGHT = 352 +_PIXEL_WIDTH = 640 +_LATENT_HEIGHT = 44 +_LATENT_WIDTH = 80 +_CHUNK_SIZE_T = 3 +_WINDOW_SIZE_T = 15 +_SINK_SIZE_T = 3 +_TEXT_TOKENS = 512 +_WARMUP_ROUNDS = 3 +_BENCHMARK_ROUNDS = 20 +_SEED = 0 + + +def _benchmark_device() -> torch.device: + """Initialize context parallelism and return this rank's GPU.""" + if int(os.environ.get("WORLD_SIZE", "1")) > 1 and not dist.is_initialized(): + init_distributed() + if dist.is_initialized(): + return torch.device("cuda", torch.cuda.current_device()) + torch.cuda.set_device(0) + return torch.device("cuda", 0) + + +def _synchronize_ranks() -> None: + """Align context-parallel ranks before a benchmark sample.""" + if dist.is_initialized(): + dist.barrier() + + +def _make_block( + config: LingbotWorldDiTNetwork14BConfig, + case: AttentionBenchmarkCase, + device: torch.device, + dtype: torch.dtype, +) -> CamCtrlBlock: + """Build a backend-selected block with shared random weights.""" + + def make(selected_backend: AttentionBackend) -> CamCtrlBlock: + return CamCtrlBlock( + dim=config.dim, + ffn_dim=config.ffn_dim, + num_heads=config.num_heads, + cross_attn_norm=config.cross_attn_norm, + eps=config.eps, + cp_method=config.cp_method, + attention_backend=selected_backend, + sdpa_backend=config.sdpa_backend, + ) + + # Allocate this 14B-sized block directly in BF16 on the target GPU. This + # changes setup memory only; initialization and checkpoint loading remain + # outside the measured region. + previous_dtype = torch.get_default_dtype() + try: + torch.set_default_dtype(dtype) + with torch.device(device): + torch.manual_seed(_SEED) + reference = make(AttentionBackend.WAN) + if case.attention_backend is AttentionBackend.WAN: + return reference + block = make(case.attention_backend) + block.load_state_dict(reference.state_dict(), strict=True) + return block + finally: + torch.set_default_dtype(previous_dtype) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason=_GPU_REASON) +@pytest.mark.parametrize( + "case", + ATTENTION_CASES, + ids=lambda case: case.pytest_id, +) +@torch.inference_mode() +def test_camctrl_dit_block_benchmark( + benchmark: BenchmarkFixture, + case: AttentionBenchmarkCase, +) -> None: + """Benchmark the CLI-resolution LingBot camera-control DiT block.""" + device = _benchmark_device() + if not torch.cuda.is_bf16_supported(): + pytest.skip("LingBot DiT block benchmark requires bfloat16 support") + + dtype = torch.bfloat16 + config = LingbotWorldDiTNetwork14BConfig( + in_dim=16 + 4 + 16, + patch_embedding_type="conv3d", + control_type="cam", + cp_method="ulysses", + attention_backend=case.attention_backend, + sdpa_backend=case.sdpa_backend, + ) + cp_size = dist.get_world_size() if dist.is_initialized() else 1 + skip_unsupported_device(case, device) + if case.attention_backend is AttentionBackend.TRITON and cp_size > 1: + pytest.skip("Triton attention does not support context parallelism") + block = _make_block(config, case, device, dtype) + block.eval() + block.update_parameters_after_loading_checkpoint() + assert block.attention_backend is case.attention_backend + assert block.sdpa_backend is case.sdpa_backend + generator = torch.Generator(device=device).manual_seed(_SEED) + + cp_group = dist.group.WORLD if cp_size > 1 else None + block.set_context_parallel_group(cp_group) + self_attention_cp_enabled = block.self_attn.is_context_parallel_enabled() + cross_attention_cp_enabled = block.cross_attn.is_context_parallel_enabled() + assert self_attention_cp_enabled == ( + case.attention_backend is AttentionBackend.WAN and cp_size > 1 + ) + assert not cross_attention_cp_enabled + + patch_t = _CHUNK_SIZE_T // config.patch_size[0] + patch_h = _LATENT_HEIGHT // config.patch_size[1] + patch_w = _LATENT_WIDTH // config.patch_size[2] + tokens_per_frame = patch_h * patch_w + global_chunk_tokens = patch_t * tokens_per_frame + global_window_tokens = _WINDOW_SIZE_T * tokens_per_frame + global_sink_tokens = _SINK_SIZE_T * tokens_per_frame + assert global_chunk_tokens % cp_size == 0 + assert global_window_tokens % cp_size == 0 + assert global_sink_tokens % cp_size == 0 + chunk_tokens = global_chunk_tokens // cp_size + window_tokens = global_window_tokens // cp_size + sink_tokens = global_sink_tokens // cp_size + head_dim = config.dim // config.num_heads + + x = torch.randn( + (chunk_tokens, config.dim), + generator=generator, + device=device, + dtype=dtype, + ) + modulation = torch.randn( + (6, config.dim), + generator=generator, + device=device, + dtype=dtype, + ) + plucker_embedding = torch.randn( + x.shape, + generator=generator, + device=device, + dtype=dtype, + ) + context = torch.randn( + (1, _TEXT_TOKENS, config.dim), + generator=generator, + device=device, + dtype=dtype, + ) + cache = block.initialize_cache( + chunk_size=chunk_tokens, + window_size=window_tokens, + sink_size=sink_tokens, + context_text=context, + ) + rope = RotaryPositionEmbedding3D( + head_dim=head_dim, + len_h=patch_h, + len_w=patch_w, + len_t=patch_t, + interleaved=True, + device=device, + ) + rope.set_context_parallel_group(cp_group) + + def forward(chunk_idx: int, rope_freqs: torch.Tensor) -> torch.Tensor: + cache.before_update(chunk_idx) + result = block( + x=x, + e=modulation, + cache=cache, + rope_freqs=rope_freqs, + plucker_embedding=plucker_embedding, + ) + cache.after_update(chunk_idx) + return result + + # Fill the fixed sink and rolling window before timing. Prepare the next + # rolling slot once, then repeatedly overwrite it to mirror denoising at + # one autoregressive position without timing cache bookkeeping. + cache_prefill_chunks = (_SINK_SIZE_T + _WINDOW_SIZE_T) // _CHUNK_SIZE_T + benchmark_ar_index = cache_prefill_chunks + rope_freqs = [rope.shift_t(idx) for idx in range(benchmark_ar_index + 1)] + for chunk_idx in range(cache_prefill_chunks): + output = forward(chunk_idx, rope_freqs[chunk_idx]) + torch.cuda.synchronize(device) + + self_attention_cp_method = ( + config.cp_method if case.attention_backend is AttentionBackend.WAN else None + ) + cross_attention_method = ( + config.cp_method if case.attention_backend is AttentionBackend.WAN else None + ) + camera_parameter_count = sum( + parameter.numel() + for name, parameter in block.named_parameters() + if name.startswith("cam_") + ) + + benchmark.group = "lingbot-camctrl-dit-block" + benchmark.extra_info.update( + { + "module": "CamCtrlBlock", + "module_owner": "lingbot", + "model_family": "lingbot-world", + "model_variant": "lingbot-world-14b", + "benchmark_scope": "whole_block_including_inherited_wan_branches", + "implementation": case.implementation, + "batch_shape": [], + "pixel_resolution": [_PIXEL_HEIGHT, _PIXEL_WIDTH], + "latent_shape": [_CHUNK_SIZE_T, _LATENT_HEIGHT, _LATENT_WIDTH], + "global_chunk_tokens": global_chunk_tokens, + "local_chunk_tokens": chunk_tokens, + "global_window_tokens": global_window_tokens, + "local_window_tokens": window_tokens, + "global_sink_tokens": global_sink_tokens, + "local_sink_tokens": sink_tokens, + "text_tokens": _TEXT_TOKENS, + "model_channels": config.dim, + "ffn_channels": config.ffn_dim, + "num_heads": config.num_heads, + "parameter_count": sum( + parameter.numel() for parameter in block.parameters() + ), + "lingbot_camera_parameter_count": camera_parameter_count, + "checkpoint": "random_init_shared_weights", + "dtype": str(dtype), + "attention_backend": case.attention_backend.value, + "sdpa_backend": case.sdpa_backend.value, + "self_attention_operator": case.self_attention_operator, + "cross_attention_operator": case.cross_attention_operator, + "projection_backend": ( + "separate_qkv" + if case.attention_backend is AttentionBackend.WAN + else "row_scaled_fp8_fused_qkv_output" + ), + "self_attention_cache_dtype": str(cache.self_attn.dtype), + "cross_attention_cache_dtype": str(cache.cross_attn.text.dtype), + "self_attention_context_parallel_method": self_attention_cp_method, + "cross_attention_method": cross_attention_method, + "context_parallel_size": cp_size, + "self_attention_context_parallel_enabled": self_attention_cp_enabled, + "cross_attention_context_parallel_enabled": (cross_attention_cp_enabled), + "distributed_sample_alignment": "barrier_before_each_round", + "cache_state": "full_sink_and_window", + "cache_prefill_chunks": cache_prefill_chunks, + "benchmark_ar_index": benchmark_ar_index, + "cache_update_bookkeeping": "excluded_from_timing", + "global_rank": dist.get_rank() if dist.is_initialized() else 0, + "gpu": torch.cuda.get_device_name(device), + "compute_capability": list(torch.cuda.get_device_capability(device)), + "torch": torch.__version__, + "cuda": torch.version.cuda, + "cudnn": torch.backends.cudnn.version(), + "warmup_rounds": _WARMUP_ROUNDS, + "benchmark_rounds": _BENCHMARK_ROUNDS, + "seed": _SEED, + } + ) + + cache.before_update(benchmark_ar_index) + torch.cuda.synchronize(device) + torch.cuda.reset_peak_memory_stats(device) + + def synchronized_forward() -> torch.Tensor: + result = block( + x=x, + e=modulation, + cache=cache, + rope_freqs=rope_freqs[benchmark_ar_index], + plucker_embedding=plucker_embedding, + ) + torch.cuda.synchronize(device) + return result + + output = benchmark.pedantic( + synchronized_forward, + setup=_synchronize_ranks, + iterations=1, + rounds=_BENCHMARK_ROUNDS, + warmup_rounds=_WARMUP_ROUNDS, + ) + cache.after_update(benchmark_ar_index) + benchmark.extra_info["peak_cuda_memory_bytes"] = torch.cuda.max_memory_allocated( + device + ) + + assert output.shape == x.shape + assert torch.isfinite(output).all() diff --git a/integrations/lingbot/benchmarks/test_network.py b/integrations/lingbot/benchmarks/test_network.py new file mode 100644 index 000000000..6e2ca9c9e --- /dev/null +++ b/integrations/lingbot/benchmarks/test_network.py @@ -0,0 +1,367 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Benchmark the complete LingBot camera-control DiT network by backend.""" + +from __future__ import annotations + +import os + +import pytest +import torch +import torch.distributed as dist +from lingbot.transformer.impl.network import ( + LingbotWorldDiTNetwork, + LingbotWorldDiTNetwork14BConfig, +) +from pytest_benchmark.fixture import BenchmarkFixture + +from flashdreams.core.attention import ContextParallelAttention +from flashdreams.core.attention.rope import RotaryPositionEmbedding3D +from flashdreams.core.distributed import init as init_distributed +from flashdreams.infra.acceleration import ( + CUDAGraphDispatch, + cuda_graph_capture_ar_index, +) +from flashdreams.infra.compile import compile_module +from flashdreams.recipes.wan.transformer.impl.modules import AttentionBackend +from integrations.lingbot.benchmarks.cases import ( + ATTENTION_CASES, + AttentionBenchmarkCase, + skip_unsupported_device, +) + +pytestmark = pytest.mark.manual + +_GPU_REASON = "LingBot DiT network benchmark requires CUDA" + +# CLI replay geometry and bounded-cache preset. The 352x640 frame becomes a +# 44x80 latent, then 2x2 DiT patching yields 22x40 tokens per +# latent frame. Each AR step contains three latent frames. +_PIXEL_HEIGHT = 352 +_PIXEL_WIDTH = 640 +_LATENT_HEIGHT = 44 +_LATENT_WIDTH = 80 +_CHUNK_SIZE_T = 3 +_WINDOW_SIZE_T = 15 +_SINK_SIZE_T = 3 +_TEXT_TOKENS = 512 +_DIFFUSION_TIMESTEP = 1000.0 +_CUDA_GRAPH_WARMUP_ITERS = 2 +_WARMUP_ROUNDS = 3 +_BENCHMARK_ROUNDS = 20 +_SEED = 0 + + +def _benchmark_device() -> torch.device: + """Initialize context parallelism and return this rank's GPU.""" + if int(os.environ.get("WORLD_SIZE", "1")) > 1 and not dist.is_initialized(): + init_distributed() + if dist.is_initialized(): + return torch.device("cuda", torch.cuda.current_device()) + torch.cuda.set_device(0) + return torch.device("cuda", 0) + + +def _synchronize_ranks() -> None: + """Align context-parallel ranks before a benchmark sample.""" + if dist.is_initialized(): + dist.barrier() + + +@pytest.mark.parametrize( + "case", + ATTENTION_CASES, + ids=lambda case: case.pytest_id, +) +@pytest.mark.skipif(not torch.cuda.is_available(), reason=_GPU_REASON) +@torch.inference_mode() +def test_dit_network_benchmark( + benchmark: BenchmarkFixture, + case: AttentionBenchmarkCase, +) -> None: + """Benchmark the compiled 14B LingBot DiT at steady state.""" + device = _benchmark_device() + if not torch.cuda.is_bf16_supported(): + pytest.skip("LingBot DiT network benchmark requires bfloat16 support") + + dtype = torch.bfloat16 + torch.manual_seed(_SEED) + skip_unsupported_device(case, device) + if ( + case.attention_backend is AttentionBackend.TRITON + and dist.is_initialized() + and dist.get_world_size() > 1 + ): + pytest.skip("Triton attention does not support context parallelism") + config = LingbotWorldDiTNetwork14BConfig( + # 16 noise channels + 4 mask channels + 16 first-frame latent + # channels before the DiT's 1x2x2 patch embedding. + in_dim=16 + 4 + 16, + patch_embedding_type="conv3d", + control_type="cam", + cp_method="ulysses", + attention_backend=case.attention_backend, + sdpa_backend=case.sdpa_backend, + ) + + # Avoid materializing the 14B random initialization as fp32 CPU weights. + previous_dtype = torch.get_default_dtype() + try: + torch.set_default_dtype(dtype) + with torch.device(device): + network = LingbotWorldDiTNetwork(config) + finally: + torch.set_default_dtype(previous_dtype) + network.eval() + network.update_parameters_after_loading_checkpoint() + parameter_count = sum(parameter.numel() for parameter in network.parameters()) + + cp_size = dist.get_world_size() if dist.is_initialized() else 1 + cp_group = dist.group.WORLD if cp_size > 1 else None + network.set_context_parallel_group(cp_group) + assert all( + block.attention_backend is case.attention_backend + and block.sdpa_backend is case.sdpa_backend + for block in network.blocks + ) + attention_modules = [ + module + for module in network.modules() + if isinstance(module, ContextParallelAttention) + ] + cudnn_attention_backends = {attention.backend for attention in attention_modules} + assert cudnn_attention_backends == ( + {"cudnn"} if case.attention_backend is AttentionBackend.WAN else set() + ) + cp_enabled_attention_modules = [ + attention + for attention in attention_modules + if attention.is_context_parallel_enabled() + ] + local_attention_methods = { + attention.method + for attention in attention_modules + if not attention.is_context_parallel_enabled() + } + assert all( + attention.context_parallel_size() == cp_size + for attention in cp_enabled_attention_modules + ) + assert all( + attention.method == config.cp_method + for attention in cp_enabled_attention_modules + ) + assert bool(cp_enabled_attention_modules) == (cp_size > 1) + + patch_t = _CHUNK_SIZE_T // config.patch_size[0] + patch_h = _LATENT_HEIGHT // config.patch_size[1] + patch_w = _LATENT_WIDTH // config.patch_size[2] + patch_volume = config.patch_size[0] * config.patch_size[1] * config.patch_size[2] + tokens_per_frame = patch_h * patch_w + global_chunk_tokens = patch_t * tokens_per_frame + global_window_tokens = _WINDOW_SIZE_T * tokens_per_frame + global_sink_tokens = _SINK_SIZE_T * tokens_per_frame + assert global_chunk_tokens % cp_size == 0 + assert global_window_tokens % cp_size == 0 + assert global_sink_tokens % cp_size == 0 + chunk_tokens = global_chunk_tokens // cp_size + window_tokens = global_window_tokens // cp_size + sink_tokens = global_sink_tokens // cp_size + head_dim = config.dim // config.num_heads + generator = torch.Generator(device=device).manual_seed(_SEED) + + x = torch.randn( + (chunk_tokens, config.in_dim * patch_volume), + generator=generator, + device=device, + dtype=dtype, + ) + control_channels = 6 if config.control_type == "cam" else 7 + plucker = torch.randn( + (chunk_tokens, control_channels * 64 * patch_volume), + generator=generator, + device=device, + dtype=dtype, + ) + timestep = torch.tensor(_DIFFUSION_TIMESTEP, device=device, dtype=dtype) + text_embeddings = torch.randn( + (1, _TEXT_TOKENS, config.text_dim), + generator=generator, + device=device, + dtype=dtype, + ) + cache = network.initialize_cache( + chunk_size=chunk_tokens, + window_size=window_tokens, + sink_size=sink_tokens, + text_embeddings=text_embeddings, + ) + rope = RotaryPositionEmbedding3D( + head_dim=head_dim, + len_h=patch_h, + len_w=patch_w, + len_t=patch_t, + interleaved=True, + device=device, + ) + rope.set_context_parallel_group(cp_group) + + network = compile_module(network) + capture_ar_index = cuda_graph_capture_ar_index( + sink_size_t=_SINK_SIZE_T, + window_size_t=_WINDOW_SIZE_T, + len_t=_CHUNK_SIZE_T, + ) + graph_dispatch = CUDAGraphDispatch( + network, + enabled=True, + capture_ar_idx=capture_ar_index, + warmup_iters=_CUDA_GRAPH_WARMUP_ITERS, + ) + + def forward(chunk_idx: int, rope_freqs: torch.Tensor) -> torch.Tensor: + return graph_dispatch.select(chunk_idx, uncond=False)( + plucker=plucker, + x=x, + timesteps=timestep, + cache=cache, + rope_freqs=rope_freqs, + current_chunk_idx=chunk_idx, + eager_mode=False, + ) + + # Fill the KV cache through the production graph threshold. At that final + # index, emulate the four scheduler evaluations that warm, capture, and + # replay the production CUDA graph before the benchmark begins. + benchmark_ar_index = capture_ar_index + 1 + rope_freqs = [ + rope.shift_t(chunk_idx) for chunk_idx in range(benchmark_ar_index + 1) + ] + for chunk_idx in range(capture_ar_index + 1): + cache.before_update(chunk_idx) + output = forward(chunk_idx, rope_freqs[chunk_idx]) + if chunk_idx == capture_ar_index: + for _ in range(_CUDA_GRAPH_WARMUP_ITERS + 1): + output = forward(chunk_idx, rope_freqs[chunk_idx]) + cache.after_update(chunk_idx) + torch.cuda.synchronize(device) + + lingbot_camera_parameter_count = sum( + parameter.numel() + for name, parameter in network.named_parameters() + if "patch_embedding_wancamctrl" in name + or "c2ws_hidden_states" in name + or ".cam_" in name + ) + benchmark.group = "lingbot-camctrl-dit-network" + benchmark.extra_info.update( + { + "network": "LingbotWorldDiTNetwork14B", + "network_owner": "lingbot", + "batch_shape": [], + "pixel_resolution": [_PIXEL_HEIGHT, _PIXEL_WIDTH], + "latent_shape": [_CHUNK_SIZE_T, _LATENT_HEIGHT, _LATENT_WIDTH], + "global_chunk_tokens": global_chunk_tokens, + "local_chunk_tokens": chunk_tokens, + "implementation": case.implementation, + "global_window_tokens": global_window_tokens, + "local_window_tokens": window_tokens, + "global_sink_tokens": global_sink_tokens, + "local_sink_tokens": sink_tokens, + "text_tokens": _TEXT_TOKENS, + "input_patch_channels": config.in_dim * patch_volume, + "plucker_patch_channels": control_channels * 64 * patch_volume, + "model_channels": config.dim, + "ffn_channels": config.ffn_dim, + "num_blocks": config.num_layers, + "num_heads": config.num_heads, + "parameter_count": parameter_count, + "lingbot_camera_parameter_count": lingbot_camera_parameter_count, + "checkpoint": "random_init", + "dtype": str(dtype), + "execution_backend": "pytorch", + "attention_backend": case.attention_backend.value, + "sdpa_backend": case.sdpa_backend.value, + "self_attention_operator": case.self_attention_operator, + "cross_attention_operator": case.cross_attention_operator, + "projection_backend": ( + "separate_qkv" + if case.attention_backend is AttentionBackend.WAN + else "row_scaled_fp8_fused_qkv_output" + ), + "self_attention_cache_dtype": str(cache[0].self_attn.dtype), + "cross_attention_cache_dtype": str(cache[0].cross_attn.text.dtype), + "self_attention_context_parallel_method": ( + config.cp_method + if case.attention_backend is AttentionBackend.WAN + else None + ), + "local_attention_methods": sorted(local_attention_methods), + "context_parallel_size": cp_size, + "context_parallel_attention_modules": len(cp_enabled_attention_modules), + "local_attention_modules": ( + len(attention_modules) - len(cp_enabled_attention_modules) + ), + "distributed_sample_alignment": "barrier_before_each_round", + "compiled": True, + "compile_mode": "max-autotune-no-cudagraphs", + "cuda_graph": True, + "cuda_graph_warmup_iters": _CUDA_GRAPH_WARMUP_ITERS, + "cache_state": "full_sink_and_window", + "cache_prefill_chunks": capture_ar_index + 1, + "benchmark_ar_index": benchmark_ar_index, + "diffusion_timestep": _DIFFUSION_TIMESTEP, + "global_rank": dist.get_rank() if dist.is_initialized() else 0, + "gpu": torch.cuda.get_device_name(device), + "compute_capability": list(torch.cuda.get_device_capability(device)), + "torch": torch.__version__, + "cuda": torch.version.cuda, + "cudnn": torch.backends.cudnn.version(), + "warmup_rounds": _WARMUP_ROUNDS, + "benchmark_rounds": _BENCHMARK_ROUNDS, + "compiler_cache_state": ( + "host-dependent; compile, autotune, and CUDA graph capture " + "excluded from measured rounds" + ), + "seed": _SEED, + } + ) + + # Scheduler evaluations at one AR position repeatedly overwrite the same + # cache slot. Cache finalization remains outside the measured callable. + cache.before_update(benchmark_ar_index) + torch.cuda.reset_peak_memory_stats(device) + + def synchronized_forward() -> torch.Tensor: + result = forward(benchmark_ar_index, rope_freqs[benchmark_ar_index]) + torch.cuda.synchronize(device) + return result + + output = benchmark.pedantic( + synchronized_forward, + setup=_synchronize_ranks, + iterations=1, + rounds=_BENCHMARK_ROUNDS, + warmup_rounds=_WARMUP_ROUNDS, + ) + cache.after_update(benchmark_ar_index) + benchmark.extra_info["peak_cuda_memory_bytes"] = torch.cuda.max_memory_allocated( + device + ) + + expected_output_shape = (chunk_tokens, config.out_dim * patch_volume) + assert output.shape == expected_output_shape + assert torch.isfinite(output).all() diff --git a/integrations/lingbot/benchmarks/test_pipeline.py b/integrations/lingbot/benchmarks/test_pipeline.py new file mode 100644 index 000000000..aeb9990fe --- /dev/null +++ b/integrations/lingbot/benchmarks/test_pipeline.py @@ -0,0 +1,582 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Steady-state full-pipeline benchmarks for LingBot streaming inference. + +Run the manual GPU benchmarks with:: + + uv run --package flashdreams-lingbot --group test pytest \ + integrations/lingbot/benchmarks/test_pipeline.py \ + -p no:manual_marker -m manual --benchmark-only -v +""" + +from __future__ import annotations + +import math +import os +from typing import Literal + +import pytest +import torch +import torch.distributed as dist +from lingbot.config import ( + PIPELINE_LINGBOT_WORLD_V2_14B_CAUSAL_FAST_TAEHV_WINDOW15_SINK3, +) +from lingbot.encoder.camctrl import CamCtrlInput, I2VCamCtrlEncoderConfig +from lingbot.pipeline import LingbotWorldInferencePipeline +from lingbot.transformer import ( + LingbotWorldTransformer, + LingbotWorldTransformerConfig, +) +from lingbot.transformer.impl.network import LingbotWorldDiTNetwork +from pytest_benchmark.fixture import BenchmarkFixture + +from flashdreams.core.attention import ContextParallelAttention +from flashdreams.core.distributed import init as init_distributed +from flashdreams.infra.config import derive_config +from flashdreams.infra.diffusion.scheduler.fm import ( + FlowMatchScheduler, + FlowMatchSchedulerConfig, +) +from flashdreams.infra.pipeline import StreamInferencePipeline +from flashdreams.recipes.taehv import TeahvVAEDecoderConfig +from flashdreams.recipes.wan.autoencoder.vae import WanVAEEncoderConfig +from flashdreams.recipes.wan.pipeline import WanInferencePipelineCache +from flashdreams.recipes.wan.transformer.impl.modules import AttentionBackend +from integrations.lingbot.benchmarks.cases import ( + ATTENTION_CASES, + AttentionBenchmarkCase, + skip_unsupported_device, +) + +pytestmark = pytest.mark.manual + +_GPU_REASON = "LingBot full-pipeline benchmark requires CUDA" + +_PIXEL_HEIGHT = 352 +_PIXEL_WIDTH = 640 +_TEXT_TOKENS = 512 +_WARMUP_ROUNDS = 3 +_BENCHMARK_ROUNDS = 20 +_SEED = 42 + + +def _benchmark_device() -> torch.device: + """Initialize context parallelism and return this rank's GPU.""" + if int(os.environ.get("WORLD_SIZE", "1")) > 1 and not dist.is_initialized(): + init_distributed() + if dist.is_initialized(): + return torch.device("cuda", torch.cuda.current_device()) + torch.cuda.set_device(0) + return torch.device("cuda", 0) + + +def _synchronize_ranks() -> None: + """Align context-parallel ranks before a benchmark sample.""" + if dist.is_initialized(): + dist.barrier() + + +def _skip_unsupported_case( + case: AttentionBenchmarkCase, + device: torch.device, +) -> None: + """Skip a case where its hardware or context-parallel contract is unmet.""" + skip_unsupported_device(case, device) + if case.attention_backend is not AttentionBackend.TRITON: + return + world_size = ( + dist.get_world_size() + if dist.is_initialized() + else int(os.environ.get("WORLD_SIZE", "1")) + ) + if world_size > 1: + pytest.skip("Triton attention does not support context parallelism") + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason=_GPU_REASON) +@pytest.mark.parametrize( + "case", + ATTENTION_CASES, + ids=lambda case: case.pytest_id, +) +def test_full_pipeline_generate_benchmark( + benchmark: BenchmarkFixture, + case: AttentionBenchmarkCase, +) -> None: + """Benchmark steady-state LingBot encode, diffuse, and decode.""" + _run_full_pipeline_benchmark(benchmark, case=case, stage="generate") + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason=_GPU_REASON) +@pytest.mark.parametrize( + "case", + ATTENTION_CASES, + ids=lambda case: case.pytest_id, +) +def test_full_pipeline_finalize_benchmark( + benchmark: BenchmarkFixture, + case: AttentionBenchmarkCase, +) -> None: + """Benchmark the LingBot DiT cache-finalization update.""" + _run_full_pipeline_benchmark(benchmark, case=case, stage="finalize") + + +@torch.inference_mode() +def _run_full_pipeline_benchmark( + benchmark: BenchmarkFixture, + *, + case: AttentionBenchmarkCase, + stage: Literal["generate", "finalize"], +) -> None: + """Run one full-pipeline lifecycle-stage benchmark.""" + device = _benchmark_device() + if not torch.cuda.is_bf16_supported(): + pytest.skip("LingBot full-pipeline benchmark requires bfloat16 support") + _skip_unsupported_case(case, device) + + torch.manual_seed(_SEED) + torch.backends.cudnn.benchmark = True + + # UMT5 is a one-shot rollout initializer, so use a correctly shaped + # precomputed embedding. The recurring pipeline remains production-like: + # LingBot camera rendering/control, four-step DiT diffusion, and TAEHV + # decoding. Initializing the 14B DiT directly on this rank's GPU avoids a + # transient fp32 CPU copy and does not affect timed steady-state stages. + pipeline_config = derive_config( + PIPELINE_LINGBOT_WORLD_V2_14B_CAUSAL_FAST_TAEHV_WINDOW15_SINK3, + name=f"lingbot-world-v2-full-pipeline-{case.implementation}-benchmark", + text_encoder=None, + enable_sync_and_profile=False, + diffusion_model={ + "seed": _SEED, + "transformer": { + "init_device": str(device), + "network": { + "attention_backend": case.attention_backend, + "sdpa_backend": case.sdpa_backend, + }, + }, + }, + ) + pipeline = pipeline_config.setup().to(device=device) + assert isinstance(pipeline, LingbotWorldInferencePipeline) + pipeline.eval() + assert pipeline.encoder is not None + assert pipeline.decoder is not None + + recurring_parameter_count = sum( + parameter.numel() for parameter in pipeline.parameters() + ) + context_parallel_attention_modules = [ + module + for module in pipeline.modules() + if isinstance(module, ContextParallelAttention) + ] + context_parallel_attention_backends = { + attention.backend for attention in context_parallel_attention_modules + } + assert context_parallel_attention_backends == ( + {"cudnn"} if case.attention_backend is AttentionBackend.WAN else set() + ) + + diffusion_config = pipeline_config.diffusion_model + transformer_config = diffusion_config.transformer + scheduler_config = diffusion_config.scheduler + encoder_config = pipeline_config.encoder + decoder_config = pipeline_config.decoder + assert isinstance(transformer_config, LingbotWorldTransformerConfig) + assert isinstance(scheduler_config, FlowMatchSchedulerConfig) + assert isinstance(encoder_config, I2VCamCtrlEncoderConfig) + assert isinstance(encoder_config.i2v.encoder, WanVAEEncoderConfig) + assert isinstance(decoder_config, TeahvVAEDecoderConfig) + assert transformer_config.network.attention_backend is case.attention_backend + assert transformer_config.network.sdpa_backend is case.sdpa_backend + + transformer = pipeline.diffusion_model.transformer + assert isinstance(transformer, LingbotWorldTransformer) + assert transformer.config is transformer_config + network = getattr(transformer.network, "_orig_mod", transformer.network) + assert isinstance(network, LingbotWorldDiTNetwork) + assert network.blocks + assert all( + block.attention_backend is case.attention_backend + and block.sdpa_backend is case.sdpa_backend + for block in network.blocks + ) + cp_size = transformer._cp_size + cp_enabled_attention_modules = [ + attention + for attention in context_parallel_attention_modules + if attention.is_context_parallel_enabled() + ] + local_attention_methods = { + attention.method + for attention in context_parallel_attention_modules + if not attention.is_context_parallel_enabled() + } + assert all( + attention.context_parallel_size() == cp_size + for attention in cp_enabled_attention_modules + ) + assert all( + attention.method == transformer_config.network.cp_method + for attention in cp_enabled_attention_modules + ) + assert bool(cp_enabled_attention_modules) == (cp_size > 1) + dtype = transformer_config.dtype + spatial_compression = int(pipeline.decoder.spatial_compression_ratio) + latent_height = _PIXEL_HEIGHT // spatial_compression + latent_width = _PIXEL_WIDTH // spatial_compression + latent_channels = int(transformer_config.network.out_dim) + text_dim = int(transformer_config.network.text_dim) + + text_embeddings = torch.zeros( + (1, _TEXT_TOKENS, text_dim), + device=device, + dtype=dtype, + ) + image = torch.zeros( + (1, 3, _PIXEL_HEIGHT, _PIXEL_WIDTH), + device=device, + dtype=dtype, + ) + + # Bypass only Wan's raw-text one-shot initializer. The base cache builder + # still creates the real recurring encoder, transformer, and decoder + # caches, and the Wan cache wrapper retains the first frame for I2V. + parent_cache = StreamInferencePipeline.initialize_cache( + pipeline, + transformer_context={ + "height": latent_height, + "width": latent_width, + "text_embeddings": text_embeddings, + "negative_text_embeddings": None, + "image_embeddings": None, + }, + ) + cache = WanInferencePipelineCache( + transformer_cache=parent_cache.transformer_cache, + encoder_cache=parent_cache.encoder_cache, + decoder_cache=parent_cache.decoder_cache, + image=image, + ) + del text_embeddings + + first_block_cache = cache.transformer_cache.network_cache.block_caches[0] + self_attention_cache_dtype = str(first_block_cache.self_attn.dtype) + cross_attention_cache_dtype = str(first_block_cache.cross_attn.text.dtype) + + first_chunk_frames = pipeline.get_num_input_frames(0) + steady_input_frames = pipeline.get_num_input_frames(1) + steady_output_frames = pipeline.get_num_output_frames(1) + + def camera_input(num_frames: int) -> CamCtrlInput: + intrinsics = torch.tensor( + [416.0, 416.0, _PIXEL_WIDTH / 2, _PIXEL_HEIGHT / 2], + device=device, + dtype=torch.float32, + ).repeat(num_frames, 1) + poses = torch.eye(4, device=device, dtype=torch.float32).repeat( + num_frames, 1, 1 + ) + return CamCtrlInput( + intrinsics=intrinsics, + poses=poses, + world_scale=1.0, + ) + + first_camera_input = camera_input(first_chunk_frames) + steady_camera_input = camera_input(steady_input_frames) + + def run_chunk( + autoregressive_index: int, + input: CamCtrlInput, + ) -> torch.Tensor: + output = pipeline.generate( + autoregressive_index=autoregressive_index, + cache=cache, + input=input, + ) + pipeline.finalize( + autoregressive_index=autoregressive_index, + cache=cache, + ) + return output + + # Fill the sink/window cache through the first CUDA-graph index. This also + # advances the reused Wan I2V encoder past its first five real VAE calls; + # steady-state timed rounds reuse its cached latent and measure only the + # LingBot camera-control work on that branch. + capture_ar_index = transformer._cuda_graph_capture_ar_idx + cache_prefill_chunks = capture_ar_index + 1 + for autoregressive_index in range(cache_prefill_chunks): + chunk_input = ( + first_camera_input if autoregressive_index == 0 else steady_camera_input + ) + output = run_chunk(autoregressive_index, chunk_input) + torch.cuda.synchronize(device) + + scheduler = pipeline.diffusion_model.scheduler + assert isinstance(scheduler, FlowMatchScheduler) + resolved_denoising_timesteps = ( + scheduler.denoising_step_list.detach().to(torch.float32).cpu().tolist() + ) + effective_dit_timesteps = ( + scheduler.denoising_step_list.detach() + .to(dtype=dtype) + .to(torch.float32) + .cpu() + .tolist() + ) + denoising_sigmas = ( + scheduler.denoising_sigmas.detach().to(torch.float32).cpu().tolist() + ) + timed_stages = ( + ["cached_i2v_and_plucker_encode", "diffuse", "taehv_decode"] + if stage == "generate" + else ["dit_cache_finalize"] + ) + untimed_lifecycle_stage = "finalize" if stage == "generate" else "generate" + measurement_start_ar_index = cache_prefill_chunks + _WARMUP_ROUNDS + measurement_end_ar_index = measurement_start_ar_index + _BENCHMARK_ROUNDS - 1 + benchmark.group = f"lingbot-full-pipeline-{stage}" + benchmark.extra_info.update( + { + "pipeline": pipeline_config.name, + "source_pipeline": ( + PIPELINE_LINGBOT_WORLD_V2_14B_CAUSAL_FAST_TAEHV_WINDOW15_SINK3.name + ), + "batch_shape": list(transformer_config.batch_shape), + "pixel_resolution": [_PIXEL_HEIGHT, _PIXEL_WIDTH], + "latent_shape": [ + transformer_config.len_t, + latent_channels, + latent_height, + latent_width, + ], + "global_chunk_tokens": ( + transformer_config.len_t + * (latent_height // transformer_config.network.patch_size[1]) + * (latent_width // transformer_config.network.patch_size[2]) + ), + "local_chunk_tokens": transformer.latent_shape[-2], + "input_frames_per_chunk": steady_input_frames, + "output_frames_per_chunk": steady_output_frames, + "text_tokens": _TEXT_TOKENS, + "text_embedding_dim": text_dim, + "num_inference_steps": scheduler_config.num_inference_steps, + "configured_denoising_timesteps": list( + scheduler_config.denoising_timesteps + ), + "resolved_denoising_timesteps_fp32": resolved_denoising_timesteps, + "effective_dit_timesteps": effective_dit_timesteps, + "denoising_sigmas_fp32": denoising_sigmas, + "context_noise": diffusion_config.context_noise, + "window_size_t": transformer_config.window_size_t, + "sink_size_t": transformer_config.sink_size_t, + "cache_prefill_chunks": cache_prefill_chunks, + "fixture_warmup_start_ar_index": cache_prefill_chunks, + "measurement_start_ar_index": measurement_start_ar_index, + "measurement_end_ar_index": measurement_end_ar_index, + "timed_stage": stage, + "timed_stages": timed_stages, + "untimed_lifecycle_stage": untimed_lifecycle_stage, + "one_shot_text_input": "synthetic_precomputed_embedding", + "first_frame_image": "zeros", + "camera_control_schedule": "repeated_static_camera", + "camera_intrinsics": [ + 416.0, + 416.0, + _PIXEL_WIDTH / 2, + _PIXEL_HEIGHT / 2, + ], + "camera_poses": "identity_4x4_repeated_per_frame", + "camera_world_scale": 1.0, + "reused_wan_i2v_vae_timed": False, + "reused_taehv_decoder_timed": stage == "generate", + "dit_checkpoint": transformer_config.checkpoint_path, + "i2v_encoder_checkpoint": encoder_config.i2v.encoder.checkpoint_path, + "decoder_checkpoint": decoder_config.checkpoint_path, + "dtype": str(dtype), + "implementation": case.implementation, + "dit_execution": "pytorch", + "configured_attention_backend": case.attention_backend.value, + "configured_sdpa_backend": case.sdpa_backend.value, + "dit_attention_backend": case.self_attention_operator, + "dit_self_attention_backend": case.self_attention_operator, + "dit_cross_attention_backend": case.cross_attention_operator, + "projection_backend": ( + "separate_qkv" + if case.attention_backend is AttentionBackend.WAN + else "row_scaled_fp8_fused_qkv_output" + ), + "dit_self_attention_kv_cache_dtype": self_attention_cache_dtype, + "dit_cross_attention_kv_cache_dtype": cross_attention_cache_dtype, + "self_attention_context_parallel_method": ( + transformer_config.network.cp_method + if case.attention_backend is AttentionBackend.WAN + else None + ), + "local_attention_methods": sorted(local_attention_methods), + "context_parallel_size": cp_size, + "context_parallel_attention_modules": len(cp_enabled_attention_modules), + "local_attention_modules": ( + len(context_parallel_attention_modules) + - len(cp_enabled_attention_modules) + ), + "distributed_sample_alignment": "barrier_before_each_round", + "dit_compiled": transformer_config.compile_network, + "dit_cuda_graph": transformer_config.use_cuda_graph, + "i2v_encoder_compiled": encoder_config.i2v.encoder.use_compile, + "i2v_encoder_cuda_graph": encoder_config.i2v.encoder.use_cuda_graph, + "decoder_compiled": decoder_config.use_compile, + "decoder_cuda_graph": decoder_config.use_cuda_graph, + "recurring_pipeline_parameter_count": recurring_parameter_count, + "global_rank": dist.get_rank() if dist.is_initialized() else 0, + "gpu": torch.cuda.get_device_name(device), + "compute_capability": list(torch.cuda.get_device_capability(device)), + "torch": torch.__version__, + "cuda": torch.version.cuda, + "cudnn": torch.backends.cudnn.version(), + "cudnn_benchmark": torch.backends.cudnn.benchmark, + "warmup_rounds": _WARMUP_ROUNDS, + "benchmark_rounds": _BENCHMARK_ROUNDS, + "startup_timing": "excluded", + "first_visible_timing": "excluded", + "compiler_cache_state": ( + "host-dependent; checkpoint loading, compile, CUDA graph " + "capture, and autotune excluded by prefill" + ), + "num_gpus_visible": torch.cuda.device_count(), + "seed": _SEED, + } + ) + + next_chunk_index = cache_prefill_chunks + latest_output: torch.Tensor | None = None + stage_peak_cuda_memory_bytes = 0 + + def record_stage_peak_memory() -> None: + nonlocal stage_peak_cuda_memory_bytes + stage_peak_cuda_memory_bytes = max( + stage_peak_cuda_memory_bytes, + int(torch.cuda.max_memory_allocated(device)), + ) + + if stage == "generate": + + def setup_generate() -> None: + _synchronize_ranks() + torch.cuda.reset_peak_memory_stats(device) + + def synchronized_generate() -> torch.Tensor: + nonlocal latest_output + latest_output = pipeline.generate( + autoregressive_index=next_chunk_index, + cache=cache, + input=steady_camera_input, + ) + torch.cuda.synchronize(device) + return latest_output + + def teardown_generate() -> None: + nonlocal next_chunk_index + record_stage_peak_memory() + pipeline.finalize( + autoregressive_index=next_chunk_index, + cache=cache, + ) + torch.cuda.synchronize(device) + next_chunk_index += 1 + + output = benchmark.pedantic( + synchronized_generate, + setup=setup_generate, + teardown=teardown_generate, + iterations=1, + rounds=_BENCHMARK_ROUNDS, + warmup_rounds=_WARMUP_ROUNDS, + ) + else: + + def setup_finalize() -> None: + nonlocal latest_output + _synchronize_ranks() + latest_output = pipeline.generate( + autoregressive_index=next_chunk_index, + cache=cache, + input=steady_camera_input, + ) + torch.cuda.synchronize(device) + _synchronize_ranks() + torch.cuda.reset_peak_memory_stats(device) + + def synchronized_finalize() -> None: + pipeline.finalize( + autoregressive_index=next_chunk_index, + cache=cache, + ) + torch.cuda.synchronize(device) + + def teardown_finalize() -> None: + nonlocal next_chunk_index + record_stage_peak_memory() + next_chunk_index += 1 + + benchmark.pedantic( + synchronized_finalize, + setup=setup_finalize, + teardown=teardown_finalize, + iterations=1, + rounds=_BENCHMARK_ROUNDS, + warmup_rounds=_WARMUP_ROUNDS, + ) + output = latest_output + + benchmark.extra_info["peak_cuda_memory_bytes"] = stage_peak_cuda_memory_bytes + assert benchmark.stats is not None + sample_times_s = benchmark.stats.stats.sorted_data + p90_index = math.ceil(0.9 * len(sample_times_s)) - 1 + median_stage_s = benchmark.stats.stats.median + p90_stage_s = sample_times_s[p90_index] + benchmark.extra_info.update( + { + f"median_{stage}_ms": median_stage_s * 1_000, + f"p90_{stage}_ms": p90_stage_s * 1_000, + f"median_{stage}_chunks_per_second": 1.0 / median_stage_s, + f"{stage}_chunks_per_second_at_p90_latency": 1.0 / p90_stage_s, + } + ) + if stage == "generate": + benchmark.extra_info.update( + { + "median_generate_only_output_fps": ( + steady_output_frames / median_stage_s + ), + "generate_only_output_fps_at_p90_latency": ( + steady_output_frames / p90_stage_s + ), + } + ) + + assert output is not None + assert output.shape == ( + steady_output_frames, + 3, + _PIXEL_HEIGHT, + _PIXEL_WIDTH, + ) + assert torch.isfinite(output).all() diff --git a/integrations/lingbot/lingbot/transformer/impl/modules.py b/integrations/lingbot/lingbot/transformer/impl/modules.py index f35f7022b..6a25d5eb7 100644 --- a/integrations/lingbot/lingbot/transformer/impl/modules.py +++ b/integrations/lingbot/lingbot/transformer/impl/modules.py @@ -23,7 +23,9 @@ import torch.nn.functional as F from torch import Tensor +from flashdreams.accelerated.multi_head_attention_triton import SDPABackend from flashdreams.recipes.wan.transformer.impl.modules import ( + AttentionBackend, Block, BlockCache, ) @@ -40,6 +42,8 @@ def __init__( cross_attn_norm: bool = True, eps: float = 1e-6, cp_method: Literal["ring", "ulysses"] = "ring", + attention_backend: AttentionBackend = AttentionBackend.TRITON, + sdpa_backend: SDPABackend = SDPABackend.TRITON, ) -> None: super().__init__( dim=dim, @@ -48,6 +52,8 @@ def __init__( cross_attn_norm=cross_attn_norm, eps=eps, cp_method=cp_method, + attention_backend=attention_backend, + sdpa_backend=sdpa_backend, ) self.cam_injector_layer1 = nn.Linear(dim, dim) self.cam_injector_layer2 = nn.Linear(dim, dim) diff --git a/integrations/lingbot/lingbot/transformer/impl/network.py b/integrations/lingbot/lingbot/transformer/impl/network.py index ea67774f8..68805bf7f 100644 --- a/integrations/lingbot/lingbot/transformer/impl/network.py +++ b/integrations/lingbot/lingbot/transformer/impl/network.py @@ -99,6 +99,8 @@ def _build_block(self, layer_idx: int) -> CamCtrlBlock: cross_attn_norm=self.cross_attn_norm, eps=self.eps, cp_method=self.cp_method, + attention_backend=self.attention_backend, + sdpa_backend=self.sdpa_backend, ) def replace_text_embeddings( diff --git a/integrations/lingbot/tests/test_transformer_cp.py b/integrations/lingbot/tests/test_transformer_cp.py index 303bc3fa8..5f4fa21dd 100644 --- a/integrations/lingbot/tests/test_transformer_cp.py +++ b/integrations/lingbot/tests/test_transformer_cp.py @@ -23,10 +23,16 @@ LingbotWorldTransformerConfig, ) from lingbot.transformer.impl.network import ( + LingbotWorldDiTNetwork, LingbotWorldDiTNetworkConfig, ) +from flashdreams.accelerated.multi_head_attention_triton import ( + SDPABackend, + TritonMultiHeadAttention, +) from flashdreams.recipes.wan.autoencoder.i2v import I2VCtrl +from flashdreams.recipes.wan.transformer.impl.modules import AttentionBackend pytestmark = pytest.mark.ci_cpu @@ -49,6 +55,7 @@ def test_lingbot_patchify_marks_i2v_and_plucker_as_patchified() -> None: compile_network=False, ) ) + assert transformer.network.attention_backend is AttentionBackend.WAN camctrl_embeddings = I2VCamCtrlEmbeddings( i2v=I2VCtrl( @@ -68,3 +75,53 @@ def test_lingbot_patchify_marks_i2v_and_plucker_as_patchified() -> None: # Idempotent once marked patchified. assert transformer.patchify_and_maybe_split_cp(patched) is patched + + +@pytest.mark.parametrize( + ("backend", "sdpa_backend"), + [ + pytest.param(AttentionBackend.WAN, SDPABackend.CUDNN, id="wan"), + pytest.param(AttentionBackend.TRITON, SDPABackend.CUDNN, id="triton-cudnn"), + pytest.param(AttentionBackend.TRITON, SDPABackend.TRITON, id="triton-fa2"), + ], +) +def test_lingbot_network_propagates_attention_backend( + backend: AttentionBackend, + sdpa_backend: SDPABackend, +) -> None: + """Propagate the configured attention backend into camera-control blocks.""" + network = LingbotWorldDiTNetwork( + LingbotWorldDiTNetworkConfig( + dim=64, + ffn_dim=128, + num_heads=4, + num_layers=1, + patch_embedding_type="linear", + control_type="cam", + attention_backend=backend, + sdpa_backend=sdpa_backend, + ) + ) + + assert network.attention_backend is backend + assert network.sdpa_backend is sdpa_backend + assert len(network.blocks) == 1 + block = network.blocks[0] + assert block.attention_backend is backend + assert block.sdpa_backend is sdpa_backend + if backend is AttentionBackend.TRITON: + assert isinstance(block.self_attn, TritonMultiHeadAttention) + assert block.self_attn.sdpa_backend is sdpa_backend + cache = block.self_attn.allocate_kv_cache( + batch_size=1, + chunk_size=2, + window_size=4, + sink_size=0, + device=torch.device("cpu"), + dtype=torch.bfloat16, + ) + assert cache.dtype is ( + torch.float8_e4m3fn + if sdpa_backend is SDPABackend.TRITON + else torch.bfloat16 + ) diff --git a/integrations/omnidreams/README.md b/integrations/omnidreams/README.md index be0561bce..b0d1828c4 100644 --- a/integrations/omnidreams/README.md +++ b/integrations/omnidreams/README.md @@ -178,6 +178,119 @@ explicitly to opt into Sparge/SageAttention-3 experiments. Use `native_dit_sparge_hybrid_period > 1` with `"sparge"` to enable the FP8 Sparge/SageAttention-3 hybrid schedule when the extension and GPU support it. +The native extension explicitly targets `12.0a` on validated compute capability +12.0 GPUs that support the architecture-specific SageAttention-3 FP4 path. On +other GPUs, including GB300, it leaves architecture selection to PyTorch and +builds SageAttention-3 stubs so its SM120a-only FP4 instructions are excluded. +Set `OMNIDREAMS_SINGLEVIEW_CUDA_ARCH_LIST` to override this behavior, or set +`TORCH_CUDA_ARCH_LIST` to use PyTorch's standard override (which takes +precedence). Explicit `12.0a` and PyTorch-default builds use separate extension +caches so an incompatible kernel image is not reused between them. + +## Run tests + +Run tests from the workspace root. Sync the OmniDreams `dev` extra, which +provides the interactive-drive test dependencies, together with the workspace +`test` group, which provides pytest and its shared plugins: + +```bash +uv sync --package flashdreams-omnidreams --extra dev --group test +``` + +Run all tests that participate in CPU or GPU CI with: + +```bash +uv run --package flashdreams-omnidreams --extra dev --group test pytest \ + integrations/omnidreams/tests \ + -m "not manual" -v +``` + +Use the tier markers to run a narrower suite: + +```bash +# CPU-safe tests +uv run --package flashdreams-omnidreams --extra dev --group test pytest \ + integrations/omnidreams/tests -m ci_cpu -v + +# Tests that require CUDA, libGL, or cv2 +uv run --package flashdreams-omnidreams --extra dev --group test pytest \ + integrations/omnidreams/tests -m ci_gpu -v +``` + +Heavy, credential-dependent, or environment-specific tests use the `manual` +marker. For example, run the end-to-end streaming pipeline test on a suitable +GPU with access to the required checkpoints: + +```bash +uv run --package flashdreams-omnidreams --extra dev --group test pytest \ + integrations/omnidreams/tests/test_omnidreams_pipeline.py::test_omnidreams_streaming_inference \ + -p no:manual_marker -m manual -v -s +``` + +The native CUDA extension build smoke test is opt-in because it performs a +clean extension build: + +```bash +OMNIDREAMS_SINGLEVIEW_RUN_NATIVE_BUILD_TEST=1 \ +uv run --package flashdreams-omnidreams --extra dev --group test pytest \ + integrations/omnidreams/tests/test_omnidreams_singleview_native.py::test_cuda_native_extension_builds \ + -m ci_gpu -v -s +``` + +Keep `--extra dev --group test` on `uv run`: it synchronizes the shared `.venv` +before launching pytest, and omitted selections may be removed. + +## Run benchmarks + +The OmniDreams benchmarks are manual, GPU-only pytest tests. Run them from the +workspace root on a supported NVIDIA GPU. First sync the OmniDreams package and +the workspace `test` dependency group, which provides both `pytest` and +`pytest-benchmark`: + +```bash +uv sync --package flashdreams-omnidreams --group test +``` + +Run the complete benchmark suite with: + +```bash +uv run --package flashdreams-omnidreams --group test pytest \ + integrations/omnidreams/benchmarks \ + -p no:manual_marker -m manual --benchmark-only -v +``` + +To run a narrower benchmark, replace the benchmark directory in that command +with one of these files: + +- `test_modules.py` benchmarks the DiT block and self-attention with the + `omnidreams_torch`, `triton_cudnn`, and `triton_fa2` implementations. + Cross-attention is unaffected by the SDPA selector and is benchmarked once + per projection backend. The backend-independent MLP is benchmarked once. +- `test_network.py` benchmarks one steady-state DiT evaluation with the + `omnidreams_torch`, `triton_cudnn`, `triton_fa2`, and native `cuda` + implementations. It uses production tensor geometry with random weights, so + checkpoint loading and startup are excluded. +- `test_pipeline.py` benchmarks steady-state generation and finalization with + the `omnidreams_torch`, `triton_cudnn`, `triton_fa2`, and native `cuda` + implementations at the runner's production 704x1280 resolution and scheduler + configuration. + +Both Triton implementations use row-scaled FP8 projections. `triton_cudnn` +uses PyTorch's cuDNN SDPA backend with a BF16 self-attention cache, while +`triton_fa2` uses Triton FlashAttention2 (FA2) with an E4M3 cache. Text and +cross-view attention remain BF16. + +The Triton cases require an NVIDIA GPU with compute capability 9.0 or newer; +they skip cleanly on older GPUs. + +Keep `--group test` on both commands. A plain +`uv sync --project integrations/omnidreams` installs only the integration's +runtime dependencies, so a later `uv run pytest` cannot find the benchmark +test tools. The benchmarks manage their warmup and measured rounds internally; +when publishing results, also record the commit, GPU and software stack, model +configuration, and any fallback warnings. + +## Run WebRTC server ## Run (shared demo API) From the repository root on a CUDA machine: diff --git a/integrations/omnidreams/benchmarks/cases.py b/integrations/omnidreams/benchmarks/cases.py new file mode 100644 index 000000000..ab62855a4 --- /dev/null +++ b/integrations/omnidreams/benchmarks/cases.py @@ -0,0 +1,197 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared OmniDreams attention benchmark cases.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Literal + +import pytest +import torch +from omnidreams.transformer.impl.modules import AttentionBackend + +from flashdreams.accelerated.multi_head_attention_triton import ( + QKVFusionOption, + SDPABackend, +) + + +@dataclass(frozen=True) +class AttentionBenchmarkCase: + """Configuration and metadata for one attention benchmark implementation.""" + + implementation: str + """Stable implementation name stored in benchmark metadata.""" + + self_attention_backend: AttentionBackend + """Self-attention implementation configured for this case.""" + + cross_attention_backend: AttentionBackend + """Cross-attention implementation configured for this case.""" + + sdpa_backend: SDPABackend + """SDPA backend configured for accelerated self-attention.""" + + self_attention_operator: str + """Self-attention operator reported in benchmark metadata.""" + + cross_attention_operator: str + """Cross-attention operator reported in benchmark metadata.""" + + use_fp8: bool = True + """Whether accelerated projections and supported attention storage use FP8.""" + + self_attn_qkv_fusion_option: QKVFusionOption = QKVFusionOption.FULL + """Projection fusion policy used by accelerated self-attention.""" + + cross_attn_qkv_fusion_option: QKVFusionOption = QKVFusionOption.FUSE_KV + """Projection fusion policy used by accelerated cross-attention.""" + + native_dit: bool = False + """Whether the full-pipeline case bypasses the PyTorch network.""" + + native_dit_backend: Literal["fp8_kvcache_cudnn", "bf16"] = "fp8_kvcache_cudnn" + """Native DiT compute backend used when ``native_dit`` is enabled.""" + + native_attention_backend: Literal["cudnn", "sparge", "sage3", "sage3_fp8"] = "cudnn" + """Native attention backend used when ``native_dit`` is enabled.""" + + minimum_compute_capability: tuple[int, int] | None = None + """Minimum CUDA compute capability; ``None`` accepts any CUDA device.""" + + @property + def pytest_id(self) -> str: + """Return the readable pytest parameter identifier.""" + return self.implementation.replace("_", "-") + + +BENCHMARK_CASES = [ + AttentionBenchmarkCase( + implementation="omnidreams_torch", + self_attention_backend=AttentionBackend.OMNIDREAMS, + cross_attention_backend=AttentionBackend.OMNIDREAMS, + sdpa_backend=SDPABackend.CUDNN, + self_attention_operator="cudnn", + cross_attention_operator="cudnn", + use_fp8=False, + self_attn_qkv_fusion_option=QKVFusionOption.NONE, + cross_attn_qkv_fusion_option=QKVFusionOption.NONE, + ), + AttentionBenchmarkCase( + implementation="triton_fa2_fp8_full", + self_attention_backend=AttentionBackend.TRITON, + cross_attention_backend=AttentionBackend.TRITON, + sdpa_backend=SDPABackend.TRITON, + self_attention_operator="triton_fa2", + cross_attention_operator="triton_fa2", + minimum_compute_capability=(9, 0), + ), + AttentionBenchmarkCase( + implementation="triton_cudnn_bf16_full", + self_attention_backend=AttentionBackend.TRITON, + cross_attention_backend=AttentionBackend.TRITON, + sdpa_backend=SDPABackend.CUDNN, + self_attention_operator="torch_cudnn_sdpa", + cross_attention_operator="torch_cudnn_sdpa", + use_fp8=False, + ), + AttentionBenchmarkCase( + implementation="triton_cudnn_bf16_full_omnidreams_cross", + self_attention_backend=AttentionBackend.TRITON, + cross_attention_backend=AttentionBackend.OMNIDREAMS, + sdpa_backend=SDPABackend.CUDNN, + self_attention_operator="torch_cudnn_sdpa", + cross_attention_operator="cudnn", + use_fp8=False, + cross_attn_qkv_fusion_option=QKVFusionOption.NONE, + ), + AttentionBenchmarkCase( + implementation="cuda", + self_attention_backend=AttentionBackend.OMNIDREAMS, + cross_attention_backend=AttentionBackend.OMNIDREAMS, + sdpa_backend=SDPABackend.CUDNN, + self_attention_operator="cudnn", + cross_attention_operator="cudnn", + use_fp8=False, + self_attn_qkv_fusion_option=QKVFusionOption.NONE, + cross_attn_qkv_fusion_option=QKVFusionOption.NONE, + native_dit=True, + ), + AttentionBenchmarkCase( + implementation="cuda_sparge", + self_attention_backend=AttentionBackend.OMNIDREAMS, + cross_attention_backend=AttentionBackend.OMNIDREAMS, + sdpa_backend=SDPABackend.CUDNN, + self_attention_operator="sparge", + cross_attention_operator="sparge", + use_fp8=False, + self_attn_qkv_fusion_option=QKVFusionOption.NONE, + cross_attn_qkv_fusion_option=QKVFusionOption.NONE, + native_dit=True, + native_attention_backend="sparge", + minimum_compute_capability=(12, 0), + ), + AttentionBenchmarkCase( + implementation="cuda_sage3", + self_attention_backend=AttentionBackend.OMNIDREAMS, + cross_attention_backend=AttentionBackend.OMNIDREAMS, + sdpa_backend=SDPABackend.CUDNN, + self_attention_operator="sage3", + cross_attention_operator="sage3", + use_fp8=False, + self_attn_qkv_fusion_option=QKVFusionOption.NONE, + cross_attn_qkv_fusion_option=QKVFusionOption.NONE, + native_dit=True, + native_dit_backend="bf16", + native_attention_backend="sage3", + minimum_compute_capability=(12, 0), + ), + AttentionBenchmarkCase( + implementation="cuda_sage3_fp8", + self_attention_backend=AttentionBackend.OMNIDREAMS, + cross_attention_backend=AttentionBackend.OMNIDREAMS, + sdpa_backend=SDPABackend.CUDNN, + self_attention_operator="sage3_fp8", + cross_attention_operator="sage3_fp8", + use_fp8=False, + self_attn_qkv_fusion_option=QKVFusionOption.NONE, + cross_attn_qkv_fusion_option=QKVFusionOption.NONE, + native_dit=True, + native_attention_backend="sage3_fp8", + minimum_compute_capability=(12, 0), + ), +] +"""Attention cases exercised by the OmniDreams benchmarks. + +Full QKV fusion only applies to self-attention because production text +cross-attention has unequal query and context widths. +""" + + +def skip_unsupported_device( + case: AttentionBenchmarkCase, + device: torch.device, +) -> None: + """Skip a benchmark case when device is older than its minimum capability.""" + minimum = case.minimum_compute_capability + if minimum is None: + return + if torch.cuda.get_device_capability(device) < minimum: + pytest.skip( + f"{case.pytest_id} attention requires compute capability " + f"{minimum[0]}.{minimum[1]}+" + ) diff --git a/integrations/omnidreams/benchmarks/test_modules.py b/integrations/omnidreams/benchmarks/test_modules.py new file mode 100644 index 000000000..9a7727b05 --- /dev/null +++ b/integrations/omnidreams/benchmarks/test_modules.py @@ -0,0 +1,554 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Microbenchmarks for Omnidreams model modules. + +Run the module benchmarks with:: + + uv run --group test pytest \ + integrations/omnidreams/benchmarks/test_modules.py \ + -p no:manual_marker -m manual --benchmark-only +""" + +from __future__ import annotations + +import pytest +import torch +from omnidreams.transformer.impl.modules import ( + AttentionBackend, + Block, +) +from omnidreams.transformer.impl.network import CosmosDiTNetworkConfig +from pytest_benchmark.fixture import BenchmarkFixture + +from flashdreams.accelerated.multi_head_attention_triton import ( + QKVFusionOption, + SDPABackend, +) +from flashdreams.core.attention.rope import RotaryPositionEmbedding3D +from integrations.omnidreams.benchmarks.cases import ( + BENCHMARK_CASES, + AttentionBenchmarkCase, + skip_unsupported_device, +) + +pytestmark = pytest.mark.manual + +_GPU_REASON = "Omnidreams DiT module benchmarks require CUDA" + +# Production single-view, 720p chunk-2 geometry. The VAE reduces 720x1280 to +# 90x160 latents, and the DiT's 2x2 spatial patching produces 45x80 tokens per +# latent frame. The local window holds three two-frame chunks. +_BATCH_SIZE = 1 +_NUM_VIEWS = 1 +_LATENT_HEIGHT = 90 +_LATENT_WIDTH = 160 +_CHUNK_SIZE_T = 2 +_WINDOW_SIZE_T = 6 +_TEXT_TOKENS = 512 +_WARMUP_ROUNDS = 5 +_BENCHMARK_ROUNDS = 50 +_SEED = 0 + +_OMNIDREAMS_TORCH_CASE = next( + case for case in BENCHMARK_CASES if case.implementation == "omnidreams_torch" +) +_HYBRID_ATTENTION_CASE = next( + case + for case in BENCHMARK_CASES + if case.implementation == "triton_cudnn_bf16_full_omnidreams_cross" +) + +_MODULE_BENCHMARK_CASES = [ + _OMNIDREAMS_TORCH_CASE, + *[ + AttentionBenchmarkCase( + implementation=( + f"triton_{'fa2' if sdpa_backend is SDPABackend.TRITON else 'cudnn'}_" + f"{'fp8' if use_fp8 else 'bf16'}_{qkv_fusion_option.value}" + ), + self_attention_backend=AttentionBackend.TRITON, + cross_attention_backend=AttentionBackend.TRITON, + sdpa_backend=sdpa_backend, + self_attention_operator=( + "triton_fa2" + if sdpa_backend is SDPABackend.TRITON + else "torch_cudnn_sdpa" + ), + cross_attention_operator=( + "triton_fa2" + if sdpa_backend is SDPABackend.TRITON + else "torch_cudnn_sdpa" + ), + use_fp8=use_fp8, + self_attn_qkv_fusion_option=qkv_fusion_option, + cross_attn_qkv_fusion_option=( + qkv_fusion_option + if qkv_fusion_option is not QKVFusionOption.FULL + else QKVFusionOption.FUSE_KV + ), + minimum_compute_capability=( + (9, 0) if use_fp8 or sdpa_backend is SDPABackend.TRITON else None + ), + ) + for sdpa_backend in SDPABackend + for use_fp8 in (False, True) + for qkv_fusion_option in QKVFusionOption + ], + _HYBRID_ATTENTION_CASE, +] +_MODULE_SELF_ATTENTION_CASES = [ + case + for case in _MODULE_BENCHMARK_CASES + if case.self_attention_backend is case.cross_attention_backend +] +_MODULE_CROSS_ATTENTION_CASES = [ + case + for case in _MODULE_BENCHMARK_CASES + if case.self_attention_backend is AttentionBackend.OMNIDREAMS + or case.self_attn_qkv_fusion_option is not QKVFusionOption.FULL +] + + +def _module_config(case: AttentionBenchmarkCase) -> CosmosDiTNetworkConfig: + """Build the network config for one module benchmark row.""" + return CosmosDiTNetworkConfig( + self_attention_backend=case.self_attention_backend, + cross_attention_backend=case.cross_attention_backend, + sdpa_backend=case.sdpa_backend, + cross_attn_sdpa_backend=case.sdpa_backend, + self_attn_qkv_fusion_option=case.self_attn_qkv_fusion_option, + cross_attn_qkv_fusion_option=case.cross_attn_qkv_fusion_option, + use_fp8=case.use_fp8, + ) + + +def _make_block( + config: CosmosDiTNetworkConfig, + case: AttentionBenchmarkCase, +) -> Block: + """Build a backend-selected block with shared random weights.""" + + def make(self_backend: AttentionBackend, cross_backend: AttentionBackend) -> Block: + # Keep this constructor in lockstep with CosmosDiTNetwork.__init__. + return Block( + x_dim=config.model_channels, + context_dim=config.crossattn_emb_channels, + num_heads=config.num_heads, + mlp_ratio=config.mlp_ratio, + use_adaln_lora=config.use_adaln_lora, + adaln_lora_dim=config.adaln_lora_dim, + enable_cross_view_attn=config.enable_cross_view_attn, + cp_method=config.cp_method, + self_attention_backend=self_backend, + cross_attention_backend=cross_backend, + sdpa_backend=config.sdpa_backend, + cross_attn_sdpa_backend=config.cross_attn_sdpa_backend, + self_attn_qkv_fusion_option=config.self_attn_qkv_fusion_option, + cross_attn_qkv_fusion_option=config.cross_attn_qkv_fusion_option, + use_fp8=config.use_fp8, + ) + + torch.manual_seed(_SEED) + omnidreams_block = make(AttentionBackend.OMNIDREAMS, AttentionBackend.OMNIDREAMS) + if ( + case.self_attention_backend is AttentionBackend.OMNIDREAMS + and case.cross_attention_backend is AttentionBackend.OMNIDREAMS + ): + return omnidreams_block + + block = make(case.self_attention_backend, case.cross_attention_backend) + block.load_state_dict(omnidreams_block.state_dict(), strict=True) + return block + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason=_GPU_REASON) +@pytest.mark.parametrize( + "case", _MODULE_BENCHMARK_CASES, ids=lambda case: case.pytest_id +) +@torch.inference_mode() +def test_dit_block_benchmark( + benchmark: BenchmarkFixture, + case: AttentionBenchmarkCase, +) -> None: + """Benchmark a production-configured DiT block with a full KV window.""" + if not torch.cuda.is_bf16_supported(): + pytest.skip("Omnidreams DiT block benchmark requires bfloat16 support") + + device = torch.device("cuda") + skip_unsupported_device(case, device) + dtype = torch.bfloat16 + config = _module_config(case) + block = _make_block(config, case).to(device=device, dtype=dtype) + block.eval() + generator = torch.Generator(device=device).manual_seed(_SEED) + + patch_t = _CHUNK_SIZE_T // config.patch_temporal + patch_h = _LATENT_HEIGHT // config.patch_spatial + patch_w = _LATENT_WIDTH // config.patch_spatial + tokens_per_frame = patch_h * patch_w + chunk_tokens = patch_t * tokens_per_frame + window_tokens = _WINDOW_SIZE_T * tokens_per_frame + head_dim = config.model_channels // config.num_heads + + x = torch.randn( + ( + _BATCH_SIZE, + _NUM_VIEWS, + patch_t, + tokens_per_frame, + config.model_channels, + ), + generator=generator, + device=device, + dtype=dtype, + ) + emb = torch.randn( + (_BATCH_SIZE, config.model_channels), + generator=generator, + device=device, + dtype=dtype, + ) + adaln_lora = torch.randn( + (_BATCH_SIZE, 3 * config.model_channels), + generator=generator, + device=device, + dtype=dtype, + ) + context = torch.randn( + ( + _BATCH_SIZE, + _NUM_VIEWS, + _TEXT_TOKENS, + config.crossattn_emb_channels, + ), + generator=generator, + device=device, + dtype=dtype, + ) + cache = block.initialize_cache( + chunk_size=chunk_tokens, + window_size=window_tokens, + sink_size=0, + context=context, + ) + rope = RotaryPositionEmbedding3D( + head_dim=head_dim, + len_h=patch_h, + len_w=patch_w, + len_t=patch_t, + h_extrapolation_ratio=3.0, + w_extrapolation_ratio=3.0, + device=device, + ) + + def forward(chunk_idx: int, rope_freqs: torch.Tensor) -> torch.Tensor: + cache.before_update(chunk_idx) + output = block( + x=x, + emb=emb, + cache=cache, + rope_freqs=rope_freqs, + adaln_lora=adaln_lora, + ) + cache.after_update(chunk_idx) + return output + + # Fill the rolling cache before timing so every measured call exercises + # steady-state attention over the full local window. Repeating the final + # chunk mirrors multiple denoising steps at one autoregressive position. + steady_chunk_idx = _WINDOW_SIZE_T // _CHUNK_SIZE_T - 1 + rope_freqs = [rope.shift_t(chunk_idx) for chunk_idx in range(steady_chunk_idx + 1)] + for chunk_idx, chunk_rope_freqs in enumerate(rope_freqs): + output = forward(chunk_idx, chunk_rope_freqs) + torch.cuda.synchronize() + + benchmark.group = "omnidreams-dit-block" + benchmark.extra_info.update( + { + "batch_size": _BATCH_SIZE, + "num_views": _NUM_VIEWS, + "latent_shape": [_CHUNK_SIZE_T, _LATENT_HEIGHT, _LATENT_WIDTH], + "chunk_tokens": chunk_tokens, + "window_tokens": window_tokens, + "text_tokens": _TEXT_TOKENS, + "model_channels": config.model_channels, + "num_heads": config.num_heads, + "parameter_count": sum( + parameter.numel() for parameter in block.parameters() + ), + "checkpoint": "random_init_shared_weights", + "dtype": str(dtype), + "implementation": case.implementation, + "sdpa_backend": case.sdpa_backend.value, + "use_fp8": case.use_fp8, + "self_attn_qkv_fusion_option": case.self_attn_qkv_fusion_option.value, + "cross_attn_qkv_fusion_option": case.cross_attn_qkv_fusion_option.value, + "attention_backend": case.self_attention_operator, + "self_attention_backend": case.self_attention_operator, + "cross_attention_backend": case.cross_attention_operator, + "self_attention_cache_dtype": str(cache.self_attn.dtype), + "cross_attention_cache_dtype": str(cache.cross_attn.dtype), + "cache_state": "full_window_static_context", + "cache_prefill_chunks": steady_chunk_idx + 1, + "benchmark_chunk_idx": steady_chunk_idx, + "gpu": torch.cuda.get_device_name(device), + "torch": torch.__version__, + "cuda": torch.version.cuda, + "cudnn": torch.backends.cudnn.version(), + "warmup_rounds": _WARMUP_ROUNDS, + "benchmark_rounds": _BENCHMARK_ROUNDS, + "seed": _SEED, + } + ) + + def synchronized_forward() -> torch.Tensor: + result = forward(steady_chunk_idx, rope_freqs[steady_chunk_idx]) + torch.cuda.synchronize() + return result + + output = benchmark.pedantic( + synchronized_forward, + iterations=1, + rounds=_BENCHMARK_ROUNDS, + warmup_rounds=_WARMUP_ROUNDS, + ) + + assert output.shape == x.shape + assert torch.isfinite(output).all() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason=_GPU_REASON) +@pytest.mark.parametrize( + "case", _MODULE_SELF_ATTENTION_CASES, ids=lambda case: case.pytest_id +) +@torch.inference_mode() +def test_self_attention_benchmark( + benchmark: BenchmarkFixture, + case: AttentionBenchmarkCase, +) -> None: + """Benchmark self-attention against a full production KV window.""" + if not torch.cuda.is_bf16_supported(): + pytest.skip("Omnidreams self-attention benchmark requires bfloat16 support") + + device = torch.device("cuda") + skip_unsupported_device(case, device) + dtype = torch.bfloat16 + config = _module_config(case) + attention = _make_block(config, case).self_attn.to(device=device, dtype=dtype) + attention.eval() + generator = torch.Generator(device=device).manual_seed(_SEED) + + patch_t = _CHUNK_SIZE_T // config.patch_temporal + patch_h = _LATENT_HEIGHT // config.patch_spatial + patch_w = _LATENT_WIDTH // config.patch_spatial + tokens_per_frame = patch_h * patch_w + chunk_tokens = patch_t * tokens_per_frame + window_tokens = _WINDOW_SIZE_T * tokens_per_frame + x = torch.randn( + ( + _BATCH_SIZE, + _NUM_VIEWS, + chunk_tokens, + config.model_channels, + ), + generator=generator, + device=device, + dtype=dtype, + ) + cache = attention.allocate_kv_cache( + batch_size=_BATCH_SIZE * _NUM_VIEWS, + chunk_size=chunk_tokens, + window_size=window_tokens, + sink_size=0, + device=device, + dtype=dtype, + ) + rope = RotaryPositionEmbedding3D( + head_dim=config.model_channels // config.num_heads, + len_h=patch_h, + len_w=patch_w, + len_t=patch_t, + h_extrapolation_ratio=3.0, + w_extrapolation_ratio=3.0, + device=device, + ) + + steady_chunk_idx = _WINDOW_SIZE_T // _CHUNK_SIZE_T - 1 + rope_freqs = [rope.shift_t(chunk_idx) for chunk_idx in range(steady_chunk_idx + 1)] + for chunk_idx, chunk_rope_freqs in enumerate(rope_freqs): + cache.before_update(chunk_idx) + output = attention(x, kv_cache=cache, rope_freqs=chunk_rope_freqs) + cache.after_update(chunk_idx) + torch.cuda.synchronize() + + benchmark.group = "omnidreams-dit-self-attention" + benchmark.extra_info.update( + { + "module": "self_attention", + "batch_size": _BATCH_SIZE, + "num_views": _NUM_VIEWS, + "latent_shape": [_CHUNK_SIZE_T, _LATENT_HEIGHT, _LATENT_WIDTH], + "chunk_tokens": chunk_tokens, + "window_tokens": window_tokens, + "model_channels": config.model_channels, + "num_heads": config.num_heads, + "parameter_count": sum( + parameter.numel() for parameter in attention.parameters() + ), + "checkpoint": "random_init_shared_weights", + "dtype": str(dtype), + "implementation": case.implementation, + "sdpa_backend": case.sdpa_backend.value, + "use_fp8": case.use_fp8, + "qkv_fusion_option": case.self_attn_qkv_fusion_option.value, + "attention_backend": case.self_attention_operator, + "cache_dtype": str(cache.dtype), + "cache_state": "full_window", + "cache_prefill_chunks": steady_chunk_idx + 1, + "benchmark_chunk_idx": steady_chunk_idx, + "gpu": torch.cuda.get_device_name(device), + "torch": torch.__version__, + "cuda": torch.version.cuda, + "cudnn": torch.backends.cudnn.version(), + "warmup_rounds": _WARMUP_ROUNDS, + "benchmark_rounds": _BENCHMARK_ROUNDS, + "seed": _SEED, + } + ) + + # Repeated denoising evaluations at one autoregressive position overwrite + # the final cache chunk while attending over the same full window. + cache.before_update(steady_chunk_idx) + + def synchronized_forward() -> torch.Tensor: + result = attention( + x, + kv_cache=cache, + rope_freqs=rope_freqs[steady_chunk_idx], + ) + torch.cuda.synchronize() + return result + + output = benchmark.pedantic( + synchronized_forward, + iterations=1, + rounds=_BENCHMARK_ROUNDS, + warmup_rounds=_WARMUP_ROUNDS, + ) + cache.after_update(steady_chunk_idx) + + assert output.shape == x.shape + assert torch.isfinite(output).all() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason=_GPU_REASON) +@pytest.mark.parametrize( + "case", _MODULE_CROSS_ATTENTION_CASES, ids=lambda case: case.pytest_id +) +@torch.inference_mode() +def test_cross_attention_benchmark( + benchmark: BenchmarkFixture, + case: AttentionBenchmarkCase, +) -> None: + """Benchmark cross-attention against the cached production text context.""" + if not torch.cuda.is_bf16_supported(): + pytest.skip("Omnidreams cross-attention benchmark requires bfloat16 support") + + device = torch.device("cuda") + skip_unsupported_device(case, device) + dtype = torch.bfloat16 + config = _module_config(case) + attention = _make_block(config, case).cross_attn.to(device=device, dtype=dtype) + attention.eval() + generator = torch.Generator(device=device).manual_seed(_SEED) + + patch_t = _CHUNK_SIZE_T // config.patch_temporal + patch_h = _LATENT_HEIGHT // config.patch_spatial + patch_w = _LATENT_WIDTH // config.patch_spatial + chunk_tokens = patch_t * patch_h * patch_w + x = torch.randn( + ( + _BATCH_SIZE, + _NUM_VIEWS, + chunk_tokens, + config.model_channels, + ), + generator=generator, + device=device, + dtype=dtype, + ) + context = torch.randn( + ( + _BATCH_SIZE, + _NUM_VIEWS, + _TEXT_TOKENS, + config.crossattn_emb_channels, + ), + generator=generator, + device=device, + dtype=dtype, + ) + cache = attention.compute_kv(context) + torch.cuda.synchronize() + + benchmark.group = "omnidreams-dit-cross-attention" + benchmark.extra_info.update( + { + "module": "cross_attention", + "batch_size": _BATCH_SIZE, + "num_views": _NUM_VIEWS, + "latent_shape": [_CHUNK_SIZE_T, _LATENT_HEIGHT, _LATENT_WIDTH], + "chunk_tokens": chunk_tokens, + "text_tokens": _TEXT_TOKENS, + "model_channels": config.model_channels, + "context_channels": config.crossattn_emb_channels, + "num_heads": config.num_heads, + "parameter_count": sum( + parameter.numel() for parameter in attention.parameters() + ), + "checkpoint": "random_init_shared_weights", + "dtype": str(dtype), + "implementation": case.implementation, + "sdpa_backend": case.sdpa_backend.value, + "use_fp8": case.use_fp8, + "qkv_fusion_option": case.cross_attn_qkv_fusion_option.value, + "attention_backend": case.cross_attention_operator, + "cache_dtype": str(cache.dtype), + "cache_state": "static_context", + "gpu": torch.cuda.get_device_name(device), + "torch": torch.__version__, + "cuda": torch.version.cuda, + "cudnn": torch.backends.cudnn.version(), + "warmup_rounds": _WARMUP_ROUNDS, + "benchmark_rounds": _BENCHMARK_ROUNDS, + "seed": _SEED, + } + ) + + def synchronized_forward() -> torch.Tensor: + result = attention(x, kv_cache=cache) + torch.cuda.synchronize() + return result + + output = benchmark.pedantic( + synchronized_forward, + iterations=1, + rounds=_BENCHMARK_ROUNDS, + warmup_rounds=_WARMUP_ROUNDS, + ) + + assert output.shape == x.shape + assert torch.isfinite(output).all() diff --git a/integrations/omnidreams/benchmarks/test_network.py b/integrations/omnidreams/benchmarks/test_network.py new file mode 100644 index 000000000..1d4bd62ff --- /dev/null +++ b/integrations/omnidreams/benchmarks/test_network.py @@ -0,0 +1,509 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Benchmark the complete Omnidreams DiT network. + +Run the benchmark with:: + + uv run --group test pytest \ + integrations/omnidreams/benchmarks/test_network.py \ + -p no:manual_marker -m manual --benchmark-only +""" + +from __future__ import annotations + +import pytest +import torch +from omnidreams.runner import DEFAULT_VIDEO_HEIGHT, DEFAULT_VIDEO_WIDTH +from omnidreams.transformer import CosmosTransformer, CosmosTransformerConfig +from omnidreams.transformer.impl.network import ( + CosmosDiTNetwork, + CosmosDiTNetworkConfig, +) +from pytest_benchmark.fixture import BenchmarkFixture + +from flashdreams.core.attention.rope import RotaryPositionEmbedding3D +from flashdreams.infra.acceleration import ( + CUDAGraphDispatch, + cuda_graph_capture_ar_index, +) +from flashdreams.infra.compile import compile_module +from integrations.omnidreams.benchmarks.cases import ( + BENCHMARK_CASES, + AttentionBenchmarkCase, + skip_unsupported_device, +) + +pytestmark = pytest.mark.manual + +_GPU_REASON = "Omnidreams DiT network benchmark requires CUDA" + +# Production single-view distilled runner geometry: 704x1280 pixels become +# 88x160 latents, the DiT consumes two latent frames per chunk, and the local +# window retains three chunks. HDMap conditioning uses 16 latent channels. +_BATCH_SIZE = 1 +_NUM_VIEWS = 1 +_PIXEL_HEIGHT = DEFAULT_VIDEO_HEIGHT +_PIXEL_WIDTH = DEFAULT_VIDEO_WIDTH +_LATENT_HEIGHT = 88 +_LATENT_WIDTH = 160 +_CHUNK_SIZE_T = 2 +_WINDOW_SIZE_T = 6 +_TEXT_TOKENS = 512 +_HDMAP_CHANNELS = 16 +_DIFFUSION_TIMESTEP = 450.0 +_WARMUP_ROUNDS = 5 +_BENCHMARK_ROUNDS = 50 +_SEED = 0 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason=_GPU_REASON) +@pytest.mark.parametrize( + "case", + [case for case in BENCHMARK_CASES if not case.native_dit], + ids=lambda case: case.pytest_id, +) +@torch.inference_mode() +def test_dit_network_benchmark( + benchmark: BenchmarkFixture, + case: AttentionBenchmarkCase, +) -> None: + """Benchmark one production compiled PyTorch DiT backend at steady state.""" + if not torch.cuda.is_bf16_supported(): + pytest.skip("Omnidreams DiT network benchmark requires bfloat16 support") + + device = torch.device("cuda") + skip_unsupported_device(case, device) + dtype = torch.bfloat16 + torch.manual_seed(_SEED) + + config = CosmosDiTNetworkConfig( + additional_concat_ch=_HDMAP_CHANNELS, + enable_cross_view_attn=False, + cp_method="ring", + self_attention_backend=case.self_attention_backend, + cross_attention_backend=case.cross_attention_backend, + sdpa_backend=case.sdpa_backend, + cross_attn_sdpa_backend=case.sdpa_backend, + self_attn_qkv_fusion_option=case.self_attn_qkv_fusion_option, + cross_attn_qkv_fusion_option=case.cross_attn_qkv_fusion_option, + use_fp8=case.use_fp8, + ) + network = CosmosDiTNetwork(config).to(device=device, dtype=dtype) + network.eval() + network.update_parameters_after_loading_checkpoint() + parameter_count = sum(parameter.numel() for parameter in network.parameters()) + assert all( + block.self_attention_backend is case.self_attention_backend + for block in network.blocks + ) + assert all( + block.cross_attention_backend is case.cross_attention_backend + for block in network.blocks + ) + assert all(block.sdpa_backend is case.sdpa_backend for block in network.blocks) + assert all( + block.cross_attn_sdpa_backend is case.sdpa_backend for block in network.blocks + ) + assert all( + block.self_attn_qkv_fusion_option is case.self_attn_qkv_fusion_option + for block in network.blocks + ) + assert all( + block.cross_attn_qkv_fusion_option is case.cross_attn_qkv_fusion_option + for block in network.blocks + ) + assert all(block.use_fp8 is case.use_fp8 for block in network.blocks) + generator = torch.Generator(device=device).manual_seed(_SEED) + + patch_t = _CHUNK_SIZE_T // config.patch_temporal + patch_h = _LATENT_HEIGHT // config.patch_spatial + patch_w = _LATENT_WIDTH // config.patch_spatial + patch_volume = config.patch_temporal * config.patch_spatial**2 + tokens_per_frame = patch_h * patch_w + chunk_tokens = patch_t * tokens_per_frame + window_tokens = _WINDOW_SIZE_T * tokens_per_frame + head_dim = config.model_channels // config.num_heads + + latent_patch_dim = config.in_channels * patch_volume + mask_patch_dim = patch_volume + hdmap_patch_dim = config.additional_concat_ch * patch_volume + x = torch.randn( + (_BATCH_SIZE, _NUM_VIEWS, chunk_tokens, latent_patch_dim), + generator=generator, + device=device, + dtype=dtype, + ) + condition_mask = torch.zeros( + (_BATCH_SIZE, _NUM_VIEWS, chunk_tokens, mask_patch_dim), + device=device, + dtype=dtype, + ) + hdmap_condition = torch.randn( + (_BATCH_SIZE, _NUM_VIEWS, chunk_tokens, hdmap_patch_dim), + generator=generator, + device=device, + dtype=dtype, + ) + timestep = torch.tensor(_DIFFUSION_TIMESTEP, device=device, dtype=dtype) + context = torch.randn( + ( + _BATCH_SIZE, + _NUM_VIEWS, + _TEXT_TOKENS, + config.crossattn_proj_in_channels, + ), + generator=generator, + device=device, + dtype=dtype, + ) + cache = network.initialize_cache( + chunk_size=chunk_tokens, + window_size=window_tokens, + sink_size=0, + context=context, + ) + rope = RotaryPositionEmbedding3D( + head_dim=head_dim, + len_h=patch_h, + len_w=patch_w, + len_t=patch_t, + h_extrapolation_ratio=3.0, + w_extrapolation_ratio=3.0, + device=device, + ) + + network = compile_module(network) + capture_chunk_idx = cuda_graph_capture_ar_index( + sink_size_t=0, + window_size_t=_WINDOW_SIZE_T, + len_t=_CHUNK_SIZE_T, + ) + graph_dispatch = CUDAGraphDispatch( + network, + enabled=True, + capture_ar_idx=capture_chunk_idx, + warmup_iters=2, + ) + + def forward(chunk_idx: int, rope_freqs: torch.Tensor) -> torch.Tensor: + return graph_dispatch.select(chunk_idx, uncond=False)( + x=x, + timesteps=timestep, + rope_freqs=rope_freqs, + cache=cache, + condition_video_input_mask=condition_mask, + current_chunk_idx=chunk_idx, + hdmap_condition=hdmap_condition, + view_indices=None, + eager_mode=False, + ) + + # Fill and roll every per-block KV cache through the production CUDA-graph + # threshold before timing. Benchmark warmups finish graph capture. + benchmark_chunk_idx = capture_chunk_idx + 1 + rope_freqs = [ + rope.shift_t(chunk_idx) for chunk_idx in range(benchmark_chunk_idx + 1) + ] + for chunk_idx in range(capture_chunk_idx + 1): + cache.before_update(chunk_idx) + output = forward(chunk_idx, rope_freqs[chunk_idx]) + cache.after_update(chunk_idx) + torch.cuda.synchronize() + + benchmark.group = "omnidreams-dit-network" + benchmark.extra_info.update( + { + "batch_size": _BATCH_SIZE, + "num_views": _NUM_VIEWS, + "pixel_resolution": [_PIXEL_HEIGHT, _PIXEL_WIDTH], + "latent_shape": [_CHUNK_SIZE_T, _LATENT_HEIGHT, _LATENT_WIDTH], + "chunk_tokens": chunk_tokens, + "window_tokens": window_tokens, + "text_tokens": _TEXT_TOKENS, + "hdmap_channels": config.additional_concat_ch, + "model_channels": config.model_channels, + "num_blocks": config.num_blocks, + "num_heads": config.num_heads, + "parameter_count": parameter_count, + "checkpoint": "random_init", + "dtype": str(dtype), + "implementation": case.implementation, + "execution_backend": "pytorch", + "sdpa_backend": case.sdpa_backend.value, + "use_fp8": case.use_fp8, + "self_attn_qkv_fusion_option": case.self_attn_qkv_fusion_option.value, + "cross_attn_qkv_fusion_option": case.cross_attn_qkv_fusion_option.value, + "attention_backend": case.self_attention_operator, + "self_attention_backend": case.self_attention_operator, + "cross_attention_backend": case.cross_attention_operator, + "self_attention_cache_dtype": str(cache.block_caches[0].self_attn.dtype), + "cross_attention_cache_dtype": str(cache.block_caches[0].cross_attn.dtype), + "compiled": True, + "cuda_graph": True, + "cache_prefill_chunks": capture_chunk_idx + 1, + "benchmark_ar_index": benchmark_chunk_idx, + "diffusion_timestep": _DIFFUSION_TIMESTEP, + "gpu": torch.cuda.get_device_name(device), + "torch": torch.__version__, + "cuda": torch.version.cuda, + "cudnn": torch.backends.cudnn.version(), + "warmup_rounds": _WARMUP_ROUNDS, + "benchmark_rounds": _BENCHMARK_ROUNDS, + "seed": _SEED, + } + ) + + # Repeated scheduler evaluations overwrite one production steady-state slot. + cache.before_update(benchmark_chunk_idx) + torch.cuda.reset_peak_memory_stats(device) + + def synchronized_forward() -> torch.Tensor: + result = forward(benchmark_chunk_idx, rope_freqs[benchmark_chunk_idx]) + torch.cuda.synchronize() + return result + + output = benchmark.pedantic( + synchronized_forward, + iterations=1, + rounds=_BENCHMARK_ROUNDS, + warmup_rounds=_WARMUP_ROUNDS, + ) + cache.after_update(benchmark_chunk_idx) + benchmark.extra_info["peak_cuda_memory_bytes"] = torch.cuda.max_memory_allocated( + device + ) + + expected_output_shape = ( + _BATCH_SIZE, + _NUM_VIEWS, + chunk_tokens, + config.out_channels * patch_volume, + ) + assert output.shape == expected_output_shape + assert torch.isfinite(output).all() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason=_GPU_REASON) +@pytest.mark.parametrize( + "case", + [case for case in BENCHMARK_CASES if case.native_dit], + ids=lambda case: case.pytest_id, +) +@torch.inference_mode() +def test_native_cuda_dit_network_benchmark( + benchmark: BenchmarkFixture, + case: AttentionBenchmarkCase, +) -> None: + """Benchmark the production native CUDA DiT backend at steady state.""" + if not torch.cuda.is_bf16_supported(): + pytest.skip("Omnidreams native DiT benchmark requires bfloat16 support") + if case.native_dit_backend == "fp8_kvcache_cudnn" and not hasattr( + torch, "float8_e4m3fn" + ): + pytest.skip("Omnidreams native DiT benchmark requires float8_e4m3fn") + + device = torch.device("cuda") + skip_unsupported_device(case, device) + dtype = torch.bfloat16 + torch.manual_seed(_SEED) + + network_config = CosmosDiTNetworkConfig( + additional_concat_ch=_HDMAP_CHANNELS, + enable_cross_view_attn=False, + cp_method="ring", + ) + transformer_config = CosmosTransformerConfig( + network=network_config, + dtype=dtype, + batch_shape=(_BATCH_SIZE,), + num_views=_NUM_VIEWS, + len_t=_CHUNK_SIZE_T, + window_size_t=_WINDOW_SIZE_T, + sink_size_t=0, + compile_network=False, + use_cuda_graph=True, + native_dit_acceleration="required", + native_dit_backend=case.native_dit_backend, + native_dit_attention_backend=case.native_attention_backend, + ) + transformer = CosmosTransformer(transformer_config).to(device=device, dtype=dtype) + transformer.eval() + parameter_count = sum(parameter.numel() for parameter in transformer.parameters()) + + x_unpatched = torch.randn( + ( + _BATCH_SIZE, + _NUM_VIEWS, + _CHUNK_SIZE_T, + network_config.in_channels, + _LATENT_HEIGHT, + _LATENT_WIDTH, + ), + device=device, + dtype=dtype, + ) + hdmap_unpatched = torch.randn( + ( + _BATCH_SIZE, + _NUM_VIEWS, + _CHUNK_SIZE_T, + network_config.additional_concat_ch, + _LATENT_HEIGHT, + _LATENT_WIDTH, + ), + device=device, + dtype=dtype, + ) + image_embeddings = torch.randn( + ( + _BATCH_SIZE, + _NUM_VIEWS, + 1, + network_config.in_channels, + _LATENT_HEIGHT, + _LATENT_WIDTH, + ), + device=device, + dtype=dtype, + ) + context = torch.randn( + ( + _BATCH_SIZE, + _NUM_VIEWS, + _TEXT_TOKENS, + network_config.crossattn_proj_in_channels, + ), + device=device, + dtype=dtype, + ) + timestep = torch.tensor(_DIFFUSION_TIMESTEP, device=device, dtype=dtype) + + cache = transformer.initialize_autoregressive_cache( + height=_LATENT_HEIGHT, + width=_LATENT_WIDTH, + text_embeddings=context, + image_embeddings=image_embeddings, + ) + x = transformer.patchify_and_maybe_split_cp(x_unpatched) + hdmap_condition = transformer.patchify_and_maybe_split_cp(hdmap_unpatched) + + patch_t = _CHUNK_SIZE_T // network_config.patch_temporal + patch_h = _LATENT_HEIGHT // network_config.patch_spatial + patch_w = _LATENT_WIDTH // network_config.patch_spatial + patch_volume = network_config.patch_temporal * network_config.patch_spatial**2 + tokens_per_frame = patch_h * patch_w + chunk_tokens = patch_t * tokens_per_frame + window_tokens = _WINDOW_SIZE_T * tokens_per_frame + + def forward() -> torch.Tensor: + return transformer.predict_flow( + noisy_latent=x, + timestep=timestep, + cache=cache, + input=hdmap_condition, + ) + + # Build the native runtime and FP8 weights, then fill and roll the cache + # through the production CUDA-graph threshold. Benchmark warmups finish + # graph capture before measured rounds. + capture_chunk_idx = transformer._cuda_graph_capture_ar_idx + benchmark_chunk_idx = capture_chunk_idx + 1 + for chunk_idx in range(capture_chunk_idx + 1): + cache.start(chunk_idx) + output = forward() + cache.finalize(chunk_idx) + torch.cuda.synchronize() + + native_selection = transformer._optimized_dit_selection + native_executor = transformer._optimized_dit_executor + assert native_selection is not None and native_selection.enabled + assert native_executor is not None + + assert native_executor._uses_fp8_dit is ( + case.native_dit_backend == "fp8_kvcache_cudnn" + ) + assert native_executor._attention_backend == case.native_attention_backend + benchmark.group = "omnidreams-dit-network" + benchmark.extra_info.update( + { + "batch_size": _BATCH_SIZE, + "num_views": _NUM_VIEWS, + "pixel_resolution": [_PIXEL_HEIGHT, _PIXEL_WIDTH], + "latent_shape": [_CHUNK_SIZE_T, _LATENT_HEIGHT, _LATENT_WIDTH], + "chunk_tokens": chunk_tokens, + "window_tokens": window_tokens, + "text_tokens": _TEXT_TOKENS, + "hdmap_channels": network_config.additional_concat_ch, + "model_channels": network_config.model_channels, + "num_blocks": network_config.num_blocks, + "num_heads": network_config.num_heads, + "parameter_count": parameter_count, + "checkpoint": "random_init", + "dtype": str(dtype), + "implementation": case.implementation, + "execution_backend": "native_cuda", + "native_dit_backend": transformer_config.native_dit_backend, + "native_dit_attention_backend": ( + transformer_config.native_dit_attention_backend + ), + "use_fp8": native_executor._uses_fp8_dit, + "attention_backend": native_executor._attention_backend, + "self_attention_backend": native_executor._attention_backend, + "cross_attention_backend": native_executor._attention_backend, + "native_extension": native_selection.reason, + "compiled": transformer_config.compile_network, + "cuda_graph": transformer_config.use_cuda_graph, + "cache_prefill_chunks": capture_chunk_idx + 1, + "benchmark_ar_index": benchmark_chunk_idx, + "diffusion_timestep": _DIFFUSION_TIMESTEP, + "gpu": torch.cuda.get_device_name(device), + "torch": torch.__version__, + "cuda": torch.version.cuda, + "cudnn": torch.backends.cudnn.version(), + "warmup_rounds": _WARMUP_ROUNDS, + "benchmark_rounds": _BENCHMARK_ROUNDS, + "seed": _SEED, + } + ) + + # Repeated scheduler evaluations overwrite one production steady-state slot. + cache.start(benchmark_chunk_idx) + torch.cuda.reset_peak_memory_stats(device) + + def synchronized_forward() -> torch.Tensor: + result = forward() + torch.cuda.synchronize() + return result + + output = benchmark.pedantic( + synchronized_forward, + iterations=1, + rounds=_BENCHMARK_ROUNDS, + warmup_rounds=_WARMUP_ROUNDS, + ) + cache.finalize(benchmark_chunk_idx) + benchmark.extra_info["peak_cuda_memory_bytes"] = torch.cuda.max_memory_allocated( + device + ) + + expected_output_shape = ( + _BATCH_SIZE, + _NUM_VIEWS, + chunk_tokens, + network_config.out_channels * patch_volume, + ) + assert output.shape == expected_output_shape + assert torch.isfinite(output).all() diff --git a/integrations/omnidreams/benchmarks/test_pipeline.py b/integrations/omnidreams/benchmarks/test_pipeline.py new file mode 100644 index 000000000..7102dbd53 --- /dev/null +++ b/integrations/omnidreams/benchmarks/test_pipeline.py @@ -0,0 +1,563 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Steady-state full-pipeline benchmark for OmniDreams streaming inference. + +Run the benchmark with:: + + uv run --group test pytest \ + integrations/omnidreams/benchmarks/test_pipeline.py \ + -p no:manual_marker -m manual --benchmark-only +""" + +from __future__ import annotations + +import math +from typing import Literal + +import pytest +import torch +from omnidreams.config import SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE_PERF +from omnidreams.pipeline import OmnidreamsPipeline +from omnidreams.runner import DEFAULT_VIDEO_HEIGHT, DEFAULT_VIDEO_WIDTH +from omnidreams.transformer import CosmosTransformer, CosmosTransformerConfig +from omnidreams.vae_native import OmnidreamsWanVAEEncoderConfig +from pytest_benchmark.fixture import BenchmarkFixture + +from flashdreams.infra.config import derive_config +from flashdreams.infra.diffusion.scheduler.fm import FlowMatchSchedulerConfig +from flashdreams.recipes.taehv import TeahvVAEDecoderConfig +from integrations.omnidreams.benchmarks.cases import ( + BENCHMARK_CASES, + AttentionBenchmarkCase, + skip_unsupported_device, +) + +pytestmark = pytest.mark.manual + +_GPU_REASON = "OmniDreams full-pipeline benchmark requires CUDA" + +_BATCH_SIZE = 1 +_NUM_VIEWS = 1 +_PIXEL_HEIGHT = DEFAULT_VIDEO_HEIGHT +_PIXEL_WIDTH = DEFAULT_VIDEO_WIDTH +_TEXT_TOKENS = 512 +_WARMUP_ROUNDS = 5 +_BENCHMARK_ROUNDS = 50 +_SEED = 0 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason=_GPU_REASON) +@pytest.mark.parametrize("case", BENCHMARK_CASES, ids=lambda case: case.pytest_id) +def test_full_pipeline_generate_benchmark( + benchmark: BenchmarkFixture, + case: AttentionBenchmarkCase, +) -> None: + """Benchmark pipeline generation for one DiT implementation.""" + _run_full_pipeline_benchmark( + benchmark, + case=case, + stage="generate", + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason=_GPU_REASON) +@pytest.mark.parametrize("case", BENCHMARK_CASES, ids=lambda case: case.pytest_id) +def test_full_pipeline_finalize_benchmark( + benchmark: BenchmarkFixture, + case: AttentionBenchmarkCase, +) -> None: + """Benchmark pipeline finalization for one DiT implementation.""" + _run_full_pipeline_benchmark( + benchmark, + case=case, + stage="finalize", + ) + + +@torch.inference_mode() +def _run_full_pipeline_benchmark( + benchmark: BenchmarkFixture, + *, + case: AttentionBenchmarkCase, + stage: Literal["generate", "finalize"], +) -> None: + """Run one DiT backend and pipeline-stage benchmark variant.""" + if not torch.cuda.is_bf16_supported(): + pytest.skip("OmniDreams full-pipeline benchmark requires bfloat16 support") + + device = torch.device("cuda") + torch.manual_seed(_SEED) + torch.backends.cudnn.benchmark = True + native_dit = case.native_dit + if ( + native_dit + and case.native_dit_backend == "fp8_kvcache_cudnn" + and not hasattr(torch, "float8_e4m3fn") + ): + pytest.skip("OmniDreams native DiT benchmark requires float8_e4m3fn") + self_attention_backend = case.self_attention_backend + skip_unsupported_device(case, device) + + # One-shot prompt and first-frame encoders run before streaming begins in + # production. Replace them with correctly shaped precomputed embeddings so + # the timed path covers the recurring HDMap encoder, diffusion, decoder, + # and cache-bookkeeping stages. + native_acceleration = "required" if native_dit else "disabled" + native_backend = case.native_dit_backend if native_dit else "bf16" + native_attention = case.native_attention_backend if native_dit else "auto" + pipeline_config = derive_config( + SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE_PERF, + name=f"omnidreams-full-pipeline-{case.pytest_id}-benchmark", + text_encoder=None, + image_encoder=None, + synthetic_text_max_length=_TEXT_TOKENS, + enable_sync_and_profile=False, + diffusion_model={ + "seed": _SEED, + "transformer": { + "compile_network": True, + "network": { + "self_attention_backend": self_attention_backend, + "cross_attention_backend": case.cross_attention_backend, + "sdpa_backend": case.sdpa_backend, + "cross_attn_sdpa_backend": case.sdpa_backend, + "self_attn_qkv_fusion_option": (case.self_attn_qkv_fusion_option), + "cross_attn_qkv_fusion_option": (case.cross_attn_qkv_fusion_option), + "use_fp8": case.use_fp8, + }, + # Keep cache finalization identical across the comparison; + # this performs the final context-noise DiT update before + # committing each autoregressive cache position. + "skip_finalize_kv_cache": False, + "native_dit_acceleration": native_acceleration, + "native_dit_backend": native_backend, + "native_dit_attention_backend": native_attention, + }, + }, + ) + pipeline = pipeline_config.setup().to(device=device) + assert isinstance(pipeline, OmnidreamsPipeline) + pipeline.eval() + assert pipeline.encoder is not None + assert pipeline.decoder is not None + + parameter_count = sum(parameter.numel() for parameter in pipeline.parameters()) + + diffusion_config = pipeline_config.diffusion_model + transformer_config = diffusion_config.transformer + scheduler_config = diffusion_config.scheduler + encoder_config = pipeline_config.encoder + decoder_config = pipeline_config.decoder + assert isinstance(transformer_config, CosmosTransformerConfig) + assert isinstance(scheduler_config, FlowMatchSchedulerConfig) + assert isinstance(encoder_config, OmnidreamsWanVAEEncoderConfig) + assert isinstance(decoder_config, TeahvVAEDecoderConfig) + network_config = transformer_config.network + assert network_config.self_attention_backend is self_attention_backend + assert network_config.cross_attention_backend is case.cross_attention_backend + assert network_config.sdpa_backend is case.sdpa_backend + assert network_config.cross_attn_sdpa_backend is case.sdpa_backend + assert ( + network_config.self_attn_qkv_fusion_option is case.self_attn_qkv_fusion_option + ) + assert ( + network_config.cross_attn_qkv_fusion_option is case.cross_attn_qkv_fusion_option + ) + assert network_config.use_fp8 is case.use_fp8 + + transformer = pipeline.diffusion_model.transformer + assert isinstance(transformer, CosmosTransformer) + assert transformer.config is transformer_config + assert transformer_config.native_dit_acceleration == native_acceleration + assert transformer_config.native_dit_backend == native_backend + assert transformer_config.native_dit_attention_backend == native_attention + assert transformer_config.skip_finalize_kv_cache is False + dtype = transformer_config.dtype + spatial_compression = int(pipeline.decoder.spatial_compression_ratio) + latent_height = _PIXEL_HEIGHT // spatial_compression + latent_width = _PIXEL_WIDTH // spatial_compression + latent_channels = int(network_config.in_channels) + text_dim = ( + int(network_config.crossattn_proj_in_channels) + if network_config.use_crossattn_projection + else int(network_config.crossattn_emb_channels) + ) + + text_embeddings = torch.zeros( + (_BATCH_SIZE, _NUM_VIEWS, _TEXT_TOKENS, text_dim), + device=device, + dtype=dtype, + ) + image_embeddings = torch.zeros( + ( + _BATCH_SIZE, + _NUM_VIEWS, + 1, + latent_channels, + latent_height, + latent_width, + ), + device=device, + dtype=dtype, + ) + cache = pipeline.initialize_cache_from_embeddings( + text_embeddings=text_embeddings, + image_embeddings=image_embeddings, + ) + del text_embeddings, image_embeddings + + first_chunk_frames = pipeline.get_num_frames(0) + steady_chunk_frames = pipeline.get_num_frames(1) + input_generator = torch.Generator(device=device).manual_seed(_SEED) + hdmap_first = ( + torch.rand( + ( + _BATCH_SIZE, + _NUM_VIEWS, + first_chunk_frames, + 3, + _PIXEL_HEIGHT, + _PIXEL_WIDTH, + ), + generator=input_generator, + device=device, + dtype=dtype, + ) + .mul_(2) + .sub_(1) + ) + hdmap_steady = ( + torch.rand( + ( + _BATCH_SIZE, + _NUM_VIEWS, + steady_chunk_frames, + 3, + _PIXEL_HEIGHT, + _PIXEL_WIDTH, + ), + generator=input_generator, + device=device, + dtype=dtype, + ) + .mul_(2) + .sub_(1) + ) + + def run_chunk(autoregressive_index: int, hdmap: torch.Tensor) -> torch.Tensor: + output = pipeline.generate( + autoregressive_index=autoregressive_index, + cache=cache, + hdmap=hdmap, + ) + pipeline.finalize(autoregressive_index=autoregressive_index, cache=cache) + return output + + # Fill the local attention window and execute the first steady-state index. + # This excludes torch.compile, CUDA-graph capture, kernel autotuning, and + # cache growth from both pytest-benchmark's warmups and measured rounds. + capture_ar_index = ( + transformer_config.sink_size_t + transformer_config.window_size_t + ) // transformer_config.len_t + cache_prefill_chunks = capture_ar_index + 1 + for autoregressive_index in range(cache_prefill_chunks): + hdmap = hdmap_first if autoregressive_index == 0 else hdmap_steady + run_chunk(autoregressive_index, hdmap) + torch.cuda.synchronize() + + native_selection = transformer._optimized_dit_selection + native_executor = transformer._optimized_dit_executor + if native_dit: + assert native_selection is not None and native_selection.enabled + assert native_executor is not None + dit_use_fp8 = native_executor._uses_fp8_dit + assert dit_use_fp8 is (case.native_dit_backend == "fp8_kvcache_cudnn") + assert native_executor._attention_backend == case.native_attention_backend + first_block_cache = cache.transformer_cache.network_cache.block_caches[0] + if dit_use_fp8: + fp8_runtime = native_executor._fp8_runtime + assert fp8_runtime is not None + for cache_name in ("k_self_fp8_caches", "v_self_fp8_caches"): + fp8_caches = fp8_runtime[cache_name] + assert fp8_caches + assert all( + cache_tensor.dtype == torch.uint8 for cache_tensor in fp8_caches + ) + dit_self_kv_cache_dtype = "float8_e4m3fn (uint8 native storage)" + if case.native_attention_backend == "sage3_fp8": + for cache_name in ( + "k_cross_sage3_fp4_caches", + "v_cross_sage3_fp4_caches", + ): + fp4_caches = fp8_runtime[cache_name] + assert fp4_caches + assert all( + cache_tensor.dtype == torch.uint8 for cache_tensor in fp4_caches + ) + for cache_name in ( + "k_cross_sage3_sf_caches", + "v_cross_sage3_sf_caches", + ): + scale_caches = fp8_runtime[cache_name] + assert scale_caches + assert all( + cache_tensor.dtype == torch.float8_e4m3fn + for cache_tensor in scale_caches + ) + dit_cross_kv_cache_dtype = "Sage3 FP4 + FP8 scale factors" + else: + for cache_name in ("k_cross_fp8_caches", "v_cross_fp8_caches"): + fp8_caches = fp8_runtime[cache_name] + assert fp8_caches + assert all( + cache_tensor.dtype == torch.uint8 for cache_tensor in fp8_caches + ) + dit_cross_kv_cache_dtype = "float8_e4m3fn (uint8 native storage)" + if case.native_attention_backend == "sparge": + dit_self_kv_cache_dtype += " + bfloat16 Sparge storage" + dit_cross_kv_cache_dtype += " + bfloat16 Sparge storage" + native_setup = "FP8 conversion, " + else: + assert native_executor._bf16_runtime is not None + dit_self_kv_cache_dtype = str(first_block_cache.self_attn.dtype) + dit_cross_kv_cache_dtype = str(first_block_cache.cross_attn.dtype) + native_setup = "" + dit_execution = "native_cuda" + dit_self_attn_qkv_fusion_option = "native_cuda" + dit_cross_attn_qkv_fusion_option = "native_cuda" + dit_attention_backend = native_executor._attention_backend + dit_self_attention_backend = native_executor._attention_backend + dit_cross_attention_backend = native_executor._attention_backend + dit_sdpa_backend = native_executor._attention_backend + dit_kv_cache_dtype = ( + f"self={dit_self_kv_cache_dtype}, cross={dit_cross_kv_cache_dtype}" + ) + native_extension = native_selection.reason + compiler_cache_state = ( + f"host-dependent; native extension build, {native_setup}CUDA graph " + "capture, and autotune excluded by prefill" + ) + else: + assert native_selection is None + assert native_executor is None + dit_execution = "pytorch" + dit_use_fp8 = case.use_fp8 + dit_self_attn_qkv_fusion_option = case.self_attn_qkv_fusion_option.value + dit_cross_attn_qkv_fusion_option = case.cross_attn_qkv_fusion_option.value + dit_attention_backend = case.self_attention_operator + dit_self_attention_backend = case.self_attention_operator + dit_cross_attention_backend = case.cross_attention_operator + dit_sdpa_backend = case.sdpa_backend.value + first_block_cache = cache.transformer_cache.network_cache.block_caches[0] + dit_self_kv_cache_dtype = str(first_block_cache.self_attn.dtype) + dit_cross_kv_cache_dtype = str(first_block_cache.cross_attn.dtype) + dit_kv_cache_dtype = ( + f"self={dit_self_kv_cache_dtype}, cross={dit_cross_kv_cache_dtype}" + ) + native_extension = None + compiler_cache_state = ( + "host-dependent; compile, CUDA graph capture, and autotune excluded " + "by prefill" + ) + + denoising_timesteps = scheduler_config.denoising_timesteps + timed_stages = ( + ["hdmap_encode", "diffuse", "decode"] if stage == "generate" else ["finalize"] + ) + untimed_lifecycle_stage = "finalize" if stage == "generate" else "generate" + benchmark.group = f"omnidreams-full-pipeline-{stage}" + benchmark.extra_info.update( + { + "pipeline": pipeline_config.name, + "source_pipeline": (SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE_PERF.name), + "batch_size": _BATCH_SIZE, + "num_views": _NUM_VIEWS, + "pixel_resolution": [_PIXEL_HEIGHT, _PIXEL_WIDTH], + "latent_shape": [ + transformer_config.len_t, + latent_channels, + latent_height, + latent_width, + ], + "frames_per_chunk": steady_chunk_frames, + "text_tokens": _TEXT_TOKENS, + "text_embedding_dim": text_dim, + "num_inference_steps": scheduler_config.num_inference_steps, + "denoising_timesteps": ( + list(denoising_timesteps) if denoising_timesteps is not None else None + ), + "context_noise": diffusion_config.context_noise, + "window_size_t": transformer_config.window_size_t, + "cache_prefill_chunks": cache_prefill_chunks, + "timed_stage": stage, + "timed_stages": timed_stages, + "untimed_lifecycle_stage": untimed_lifecycle_stage, + "one_shot_inputs": "synthetic_precomputed_embeddings", + "dit_checkpoint": transformer_config.checkpoint_path, + "hdmap_encoder_checkpoint": encoder_config.checkpoint_path, + "decoder_checkpoint": decoder_config.checkpoint_path, + "dtype": str(dtype), + "implementation": case.implementation, + "dit_execution": dit_execution, + "dit_sdpa_backend": dit_sdpa_backend, + "dit_use_fp8": dit_use_fp8, + "dit_self_attn_qkv_fusion_option": dit_self_attn_qkv_fusion_option, + "dit_cross_attn_qkv_fusion_option": dit_cross_attn_qkv_fusion_option, + "dit_attention_backend": dit_attention_backend, + "dit_self_attention_backend": dit_self_attention_backend, + "dit_cross_attention_backend": dit_cross_attention_backend, + "dit_kv_cache_dtype": dit_kv_cache_dtype, + "dit_self_attention_kv_cache_dtype": dit_self_kv_cache_dtype, + "dit_cross_attention_kv_cache_dtype": dit_cross_kv_cache_dtype, + "native_dit_acceleration": transformer_config.native_dit_acceleration, + "native_dit_backend": transformer_config.native_dit_backend, + "native_dit_attention_backend": ( + transformer_config.native_dit_attention_backend + ), + "native_extension": native_extension, + "dit_compiled": transformer_config.compile_network, + "dit_cuda_graph": transformer_config.use_cuda_graph, + "skip_finalize_kv_cache": transformer_config.skip_finalize_kv_cache, + "hdmap_encoder_compiled": encoder_config.use_compile, + "hdmap_encoder_cuda_graph": encoder_config.use_cuda_graph, + "decoder_compiled": decoder_config.use_compile, + "decoder_cuda_graph": decoder_config.use_cuda_graph, + "parameter_count": parameter_count, + "gpu": torch.cuda.get_device_name(device), + "torch": torch.__version__, + "cuda": torch.version.cuda, + "cudnn": torch.backends.cudnn.version(), + "cudnn_benchmark": torch.backends.cudnn.benchmark, + "warmup_rounds": _WARMUP_ROUNDS, + "benchmark_rounds": _BENCHMARK_ROUNDS, + "startup_timing": "excluded", + "first_visible_timing": "excluded", + "compiler_cache_state": compiler_cache_state, + "num_gpus": torch.cuda.device_count(), + "seed": _SEED, + } + ) + + next_chunk_index = cache_prefill_chunks + latest_output: torch.Tensor | None = None + stage_peak_cuda_memory_bytes = 0 + + def record_stage_peak_memory() -> None: + nonlocal stage_peak_cuda_memory_bytes + stage_peak_cuda_memory_bytes = max( + stage_peak_cuda_memory_bytes, + int(torch.cuda.max_memory_allocated(device)), + ) + + if stage == "generate": + + def setup_generate() -> None: + torch.cuda.reset_peak_memory_stats(device) + + def synchronized_generate() -> torch.Tensor: + nonlocal latest_output + latest_output = pipeline.generate( + autoregressive_index=next_chunk_index, + cache=cache, + hdmap=hdmap_steady, + ) + torch.cuda.synchronize() + return latest_output + + def teardown_generate() -> None: + nonlocal next_chunk_index + record_stage_peak_memory() + pipeline.finalize( + autoregressive_index=next_chunk_index, + cache=cache, + ) + torch.cuda.synchronize() + next_chunk_index += 1 + + output = benchmark.pedantic( + synchronized_generate, + setup=setup_generate, + teardown=teardown_generate, + iterations=1, + rounds=_BENCHMARK_ROUNDS, + warmup_rounds=_WARMUP_ROUNDS, + ) + else: + + def setup_finalize() -> None: + nonlocal latest_output + latest_output = pipeline.generate( + autoregressive_index=next_chunk_index, + cache=cache, + hdmap=hdmap_steady, + ) + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats(device) + + def synchronized_finalize() -> None: + pipeline.finalize( + autoregressive_index=next_chunk_index, + cache=cache, + ) + torch.cuda.synchronize() + + def teardown_finalize() -> None: + nonlocal next_chunk_index + record_stage_peak_memory() + next_chunk_index += 1 + + benchmark.pedantic( + synchronized_finalize, + setup=setup_finalize, + teardown=teardown_finalize, + iterations=1, + rounds=_BENCHMARK_ROUNDS, + warmup_rounds=_WARMUP_ROUNDS, + ) + output = latest_output + + benchmark.extra_info["peak_cuda_memory_bytes"] = stage_peak_cuda_memory_bytes + assert benchmark.stats is not None + sample_times_s = benchmark.stats.stats.sorted_data + p90_index = math.ceil(0.9 * len(sample_times_s)) - 1 + median_stage_s = benchmark.stats.stats.median + p90_stage_s = sample_times_s[p90_index] + benchmark.extra_info.update( + { + f"median_{stage}_ms": median_stage_s * 1_000, + f"p90_{stage}_ms": p90_stage_s * 1_000, + "median_chunks_per_second": 1.0 / median_stage_s, + "p90_chunks_per_second": 1.0 / p90_stage_s, + } + ) + if stage == "generate": + benchmark.extra_info.update( + { + "median_output_fps": steady_chunk_frames / median_stage_s, + "p90_output_fps": steady_chunk_frames / p90_stage_s, + } + ) + + assert output is not None + assert output.shape == ( + _BATCH_SIZE, + _NUM_VIEWS, + steady_chunk_frames, + 3, + _PIXEL_HEIGHT, + _PIXEL_WIDTH, + ) + assert torch.isfinite(output).all() diff --git a/integrations/omnidreams/omnidreams/config.py b/integrations/omnidreams/omnidreams/config.py index 9031b4bca..af566ad51 100644 --- a/integrations/omnidreams/omnidreams/config.py +++ b/integrations/omnidreams/omnidreams/config.py @@ -153,6 +153,43 @@ def _lightvae_fp8_state_path() -> str | None: """Performance-tuned variant: enable ``use_compile`` / ``use_cuda_graph`` on the image encoder, the per-AR-step encoder, and the decoder.""" +SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE_CUDA_CUDNN = cast( + OmnidreamsPipelineConfig, + derive_config( + SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE_PERF, + name="omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae-cuda-cudnn", + diffusion_model=dict( + transformer=dict( + native_dit_acceleration="required", + native_dit_backend="fp8_kvcache_cudnn", + native_dit_attention_backend="cudnn", + ), + ), + ), +) # ty:ignore[redundant-cast] + +SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE_CUDA_SPARGE = cast( + OmnidreamsPipelineConfig, + derive_config( + SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE_CUDA_CUDNN, + name="omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae-cuda-sparge", + diffusion_model=dict( + transformer=dict(native_dit_attention_backend="sparge"), + ), + ), +) # ty:ignore[redundant-cast] + +SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE_CUDA_SAGE3_FP8 = cast( + OmnidreamsPipelineConfig, + derive_config( + SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE_CUDA_CUDNN, + name="omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae-cuda-sage3-fp8", + diffusion_model=dict( + transformer=dict(native_dit_attention_backend="sage3_fp8"), + ), + ), +) # ty:ignore[redundant-cast] + SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE_NATIVE_PERF = cast( OmnidreamsPipelineConfig, derive_config( @@ -412,6 +449,9 @@ def _lightvae_fp8_state_path() -> str | None: for cfg in ( SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE, SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE_PERF, + SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE_CUDA_CUDNN, + SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE_CUDA_SPARGE, + SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE_CUDA_SAGE3_FP8, SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE_NATIVE_PERF, SV_2STEPS_CHUNK2_LOC6_VAE_VAE, SV_2STEPS_CHUNK3_LOC6_VAE_VAE, @@ -465,6 +505,33 @@ def _lightvae_fp8_state_path() -> str | None: prompt=_DEFAULT_PROMPT_1V, ) +RUNNER_SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE_CUDA_CUDNN = ( + OmnidreamsRunnerConfig( + runner_name="omnidreams-cuda-cudnn", + description="Single-view chunk2 native CUDA DiT with cuDNN attention.", + pipeline=SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE_CUDA_CUDNN, + prompt=_DEFAULT_PROMPT_1V, + ) +) + +RUNNER_SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE_CUDA_SPARGE = ( + OmnidreamsRunnerConfig( + runner_name="omnidreams-cuda-sparge", + description="Single-view chunk2 native CUDA DiT with Sparge attention.", + pipeline=SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE_CUDA_SPARGE, + prompt=_DEFAULT_PROMPT_1V, + ) +) + +RUNNER_SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE_CUDA_SAGE3_FP8 = ( + OmnidreamsRunnerConfig( + runner_name="omnidreams-cuda-sage3fp8", + description="Single-view chunk2 native CUDA DiT with SageAttention-3 FP8.", + pipeline=SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE_CUDA_SAGE3_FP8, + prompt=_DEFAULT_PROMPT_1V, + ) +) + RUNNER_SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE_NATIVE_PERF = OmnidreamsRunnerConfig( runner_name=SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE_NATIVE_PERF.name, description=( @@ -569,6 +636,9 @@ def _lightvae_fp8_state_path() -> str | None: for cfg in ( RUNNER_SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE, RUNNER_SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE_PERF, + RUNNER_SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE_CUDA_CUDNN, + RUNNER_SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE_CUDA_SPARGE, + RUNNER_SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE_CUDA_SAGE3_FP8, ) } """All shipped Omnidreams runners (single- and multi-view variants), diff --git a/integrations/omnidreams/omnidreams/native/omnidreams_singleview.py b/integrations/omnidreams/omnidreams/native/omnidreams_singleview.py index 0e9607e84..c0d082370 100644 --- a/integrations/omnidreams/omnidreams/native/omnidreams_singleview.py +++ b/integrations/omnidreams/omnidreams/native/omnidreams_singleview.py @@ -55,10 +55,17 @@ _NATIVE_CUDA_ARCH_LIST_ENV = "OMNIDREAMS_SINGLEVIEW_CUDA_ARCH_LIST" _DISABLE_SAGE3_ENV = "OMNIDREAMS_SINGLEVIEW_DISABLE_SAGE3" _PYTORCH_CUDA_ARCH_LIST_ENV = "TORCH_CUDA_ARCH_LIST" -_DEFAULT_CUDA_ARCH_LIST = "12.0a" +_PYTORCH_DEFAULT_CUDA_ARCH_LIST = "pytorch-default" +# CUDA reports capability 12.0 without the "a" suffix, so mirror the +# conservative device allowlist used by sage3_is_runtime_supported(). +_SM120A_DEVICE_NAME_MARKERS = ( + "GeForce RTX 5090", + "RTX PRO 6000", + "RTX 6000", +) _native_build_module: ModuleType | None = None -_extension: dict[bool, ModuleType] = {} +_extension: dict[tuple[bool, str], ModuleType] = {} _extension_load_error: Exception | None = None _state_lock = threading.RLock() _dll_directory_handles: list[object] = [] @@ -366,8 +373,12 @@ def _file_sha256(path: Path) -> str: return digest.hexdigest() -def _sage3_disabled() -> bool: - return os.environ.get(_DISABLE_SAGE3_ENV, "").strip().lower() in {"1", "true"} +def _sage3_disabled(cuda_arch_list: str | None = None) -> bool: + if os.environ.get(_DISABLE_SAGE3_ENV, "").strip().lower() in {"1", "true"}: + return True + if cuda_arch_list is None: + cuda_arch_list = _effective_cuda_arch_list() + return cuda_arch_list != "12.0a" def _extension_sources() -> list[Path]: @@ -433,12 +444,20 @@ def _source_fingerprint() -> str: return digest.hexdigest() -def _extension_name(thirdparty_info: dict[str, Any]) -> str: - has_sage3 = int(not _sage3_disabled()) +def _extension_name( + thirdparty_info: dict[str, Any], + *, + cuda_arch_list: str | None = None, +) -> str: + cuda_arch_list = _cuda_arch_identity( + _effective_cuda_arch_list() if cuda_arch_list is None else cuda_arch_list + ) + has_sage3 = int(not _sage3_disabled(cuda_arch_list)) digest = hashlib.sha256() digest.update(_source_fingerprint().encode("ascii")) digest.update(json.dumps(thirdparty_info, sort_keys=True).encode("utf-8")) digest.update(f"sage3={has_sage3}".encode("ascii")) + digest.update(f"cuda_arch_list={cuda_arch_list}".encode("ascii")) return f"omnidreams_singleview_native_sage3_{has_sage3}_{digest.hexdigest()[:12]}" @@ -463,19 +482,34 @@ def _resolved_max_jobs(max_jobs: int | str | None) -> str | None: return str(min(os.cpu_count() or 1, _DEFAULT_MAX_JOBS_CAP)) -def _resolved_cuda_arch_list() -> str | None: - if os.environ.get(_PYTORCH_CUDA_ARCH_LIST_ENV): +def _detected_cuda_arch_list() -> str | None: + try: + import torch + + if not torch.cuda.is_available(): + return None + if torch.cuda.get_device_capability() != (12, 0): + return None + device_name = torch.cuda.get_device_name() + if not any(marker in device_name for marker in _SM120A_DEVICE_NAME_MARKERS): + return None + return "12.0a" + except Exception: return None - return os.environ.get(_NATIVE_CUDA_ARCH_LIST_ENV, _DEFAULT_CUDA_ARCH_LIST) -def _effective_cuda_arch_list() -> str: - return os.environ.get( - _PYTORCH_CUDA_ARCH_LIST_ENV, - os.environ.get(_NATIVE_CUDA_ARCH_LIST_ENV, _DEFAULT_CUDA_ARCH_LIST), +def _effective_cuda_arch_list() -> str | None: + return ( + os.environ.get(_PYTORCH_CUDA_ARCH_LIST_ENV) + or os.environ.get(_NATIVE_CUDA_ARCH_LIST_ENV) + or _detected_cuda_arch_list() ) +def _cuda_arch_identity(cuda_arch_list: str | None) -> str: + return cuda_arch_list or _PYTORCH_DEFAULT_CUDA_ARCH_LIST + + def _python_package_dir(package: str) -> Path | None: spec = importlib.util.find_spec(package) if spec is None or spec.submodule_search_locations is None: @@ -505,14 +539,13 @@ def _scoped_torch_max_jobs(max_jobs: int | str | None) -> Iterator[None]: @contextlib.contextmanager -def _scoped_cuda_arch_list() -> Iterator[None]: - resolved = _resolved_cuda_arch_list() - if resolved is None: +def _scoped_cuda_arch_list(cuda_arch_list: str | None) -> Iterator[None]: + if cuda_arch_list is None: yield return previous = os.environ.get(_PYTORCH_CUDA_ARCH_LIST_ENV) - os.environ[_PYTORCH_CUDA_ARCH_LIST_ENV] = resolved + os.environ[_PYTORCH_CUDA_ARCH_LIST_ENV] = cuda_arch_list try: yield finally: @@ -540,8 +573,11 @@ def load_extension( global _extension, _extension_load_error with _state_lock: - sage3_disabled = _sage3_disabled() - if (extension := _extension.get(sage3_disabled)) is not None: + cuda_arch_list = _effective_cuda_arch_list() + cuda_arch_identity = _cuda_arch_identity(cuda_arch_list) + sage3_disabled = _sage3_disabled(cuda_arch_identity) + extension_key = (sage3_disabled, cuda_arch_identity) + if (extension := _extension.get(extension_key)) is not None: return extension _extension_load_error = None @@ -551,7 +587,10 @@ def load_extension( from torch.utils.cpp_extension import load as load_torch_extension thirdparty_info = validate_thirdparty() - extension_name = _extension_name(thirdparty_info) + extension_name = _extension_name( + thirdparty_info, + cuda_arch_list=cuda_arch_identity, + ) has_sage3 = int(not sage3_disabled) cutlass_dir = Path(thirdparty_info["cutlass"]["path"]) cutlass_include = cutlass_dir / "include" @@ -570,8 +609,11 @@ def load_extension( extension_build_dir.mkdir(parents=True, exist_ok=True) _add_windows_cuda_dll_directories(cudnn_package_dir) - with _scoped_torch_max_jobs(max_jobs), _scoped_cuda_arch_list(): - _extension[sage3_disabled] = load_torch_extension( + with ( + _scoped_torch_max_jobs(max_jobs), + _scoped_cuda_arch_list(cuda_arch_list), + ): + _extension[extension_key] = load_torch_extension( name=extension_name, sources=[str(source) for source in _extension_sources()], build_directory=str(extension_build_dir), @@ -636,7 +678,7 @@ def load_extension( "-DOMNIDREAMS_SINGLEVIEW_SPARGE_ATTN_SHA=" f'\\"{thirdparty_info["SpargeAttn"]["commit"]}\\"', "-DOMNIDREAMS_SINGLEVIEW_CUDA_ARCH_LIST=" - f'\\"{_effective_cuda_arch_list()}\\"', + f'\\"{cuda_arch_identity}\\"', ], extra_cuda_cflags=[ # Assume MSVC for Windows @@ -677,7 +719,7 @@ def load_extension( except Exception as exc: # pragma: no cover - environment-specific build path _extension_load_error = exc return None - return _extension[sage3_disabled] + return _extension[extension_key] def extension_load_error() -> Exception | None: diff --git a/integrations/omnidreams/omnidreams/transformer/impl/modules.py b/integrations/omnidreams/omnidreams/transformer/impl/modules.py index ab4259a3b..3bea300b9 100644 --- a/integrations/omnidreams/omnidreams/transformer/impl/modules.py +++ b/integrations/omnidreams/omnidreams/transformer/impl/modules.py @@ -17,18 +17,36 @@ import math from dataclasses import dataclass +from enum import Enum from typing import Literal +import nvtx import torch import torch.nn as nn from einops import rearrange, repeat from torch import Tensor from torch.distributed import ProcessGroup +from flashdreams.accelerated.multi_head_attention import AttentionType, QKNormScope +from flashdreams.accelerated.multi_head_attention_triton import ( + QKVFusionOption, + SDPABackend, + TritonMultiHeadAttention, +) from flashdreams.core.attention import BlockKVCache, ContextParallelAttention from flashdreams.core.attention.rope import apply_rope_freqs +class AttentionBackend(str, Enum): + """Attention implementation used by an Omnidreams DiT block.""" + + OMNIDREAMS = "omnidreams" + """Use the integration's context-parallel cuDNN attention.""" + + TRITON = "triton" + """Use Triton-accelerated attention for the selected branch.""" + + class GPT2FeedForward(nn.Module): """GPT-2 style feed-forward network with GELU activation.""" @@ -282,6 +300,7 @@ def set_context_parallel_group(self, cp_group: ProcessGroup | None) -> None: self.attn_op.set_context_parallel_group(cp_group=cp_group) def is_context_parallel_enabled(self) -> bool: + """Whether context parallelism is active for attention.""" return self.attn_op.is_context_parallel_enabled() def context_parallel_size(self) -> int: @@ -306,7 +325,7 @@ def _compute_or_update_kv_cache( """ batch_shape = context.shape[:-2] batch_size = math.prod(batch_shape) - L = context.shape[-2] + L, D = context.shape[-2:] n, d = self.n_heads, self.head_dim k = self.k_norm(self.k_proj(context).reshape(batch_size, L, n, d)) @@ -337,7 +356,7 @@ def update_kv( """Append K/V computed from ``x`` into an existing ``kv_cache``.""" return self._compute_or_update_kv_cache(x, kv_cache, rope_freqs) - def apply_kv( + def query_kv( self, x: Tensor, kv_cache: BlockKVCache, @@ -390,13 +409,13 @@ def forward( """ if update_kv_cache: kv_cache = self.update_kv(x, kv_cache, rope_freqs) - return self.apply_kv(x, kv_cache, rope_freqs) + return self.query_kv(x, kv_cache, rope_freqs) class SelfAttention(MultiHeadAttention): """Self-attention: queries and K/V are derived from the same ``x`` each step.""" - def initialize_cache( + def allocate_kv_cache( self, batch_size: int, chunk_size: int, @@ -405,7 +424,7 @@ def initialize_cache( device: torch.device, dtype: torch.dtype, ) -> BlockKVCache: - """Initialize KV cache for streaming self-attention. + """Allocate a KV cache for streaming self-attention. Args: batch_size: Flattened batch size used by attention. @@ -443,14 +462,6 @@ def forward( class CrossAttention(MultiHeadAttention): """Cross-attention: K/V live only in ``kv_cache``; ``forward`` does not refresh them.""" - def initialize_cache( - self, - context: Tensor, # [B, V, L, D] - ) -> BlockKVCache: - """Initialize cross-attention cache from the provided context.""" - cache = self.compute_kv(context) - return cache - def forward( self, x: Tensor, @@ -460,6 +471,222 @@ def forward( return super().forward(x, kv_cache, rope_freqs=None, update_kv_cache=False) +class TritonCrossAttention(TritonMultiHeadAttention): + """Static-context cross-attention backed by TMA FlashAttention2.""" + + @property + def query_projection(self) -> nn.Linear: + """Return the canonical query projection.""" + return self.q_proj + + @property + def key_projection(self) -> nn.Linear: + """Return the canonical key projection.""" + return self.k_proj + + @property + def value_projection(self) -> nn.Linear: + """Return the canonical value projection.""" + return self.v_proj + + @property + def output_projection(self) -> nn.Linear: + """Return the canonical output projection.""" + return self.output_proj + + @property + def query_norm(self) -> nn.Module: + """Return the canonical query normalization.""" + return self.q_norm + + @property + def key_norm(self) -> nn.Module: + """Return the canonical key normalization.""" + return self.k_norm + + def __init__( + self, + query_dim: int, + context_dim: int | None = None, + n_heads: int = 8, + head_dim: int = 64, + cp_method: Literal["ring", "ulysses"] = "ring", + sdpa_backend: SDPABackend = SDPABackend.TRITON, + qkv_fusion_option: QKVFusionOption = QKVFusionOption.FUSE_KV, + use_fp8: bool = True, + ) -> None: + """Initialize bias-free Triton cross-attention. + + Args: + query_dim: Feature dimension of query tokens and projected output. + context_dim: Feature dimension of key/value tokens. ``None`` uses + ``query_dim``. + n_heads: Number of attention heads. + head_dim: Per-head feature dimension. + cp_method: Ignored context-parallel method retained for constructor + compatibility with Omnidreams attention. + sdpa_backend: Scaled-dot-product attention implementation. + qkv_fusion_option: Projection fusion policy. + use_fp8: Whether projection GEMMs and supported attention storage + use FP8. + """ + del cp_method + super().__init__( + query_dim=query_dim, + context_dim=context_dim, + n_heads=n_heads, + attention_type=AttentionType.CROSS_ATTENTION, + head_dim=head_dim, + qkv_fusion_option=qkv_fusion_option, + qk_norm_eps=1e-6, + qk_norm_scope=QKNormScope.HEAD, + rope_interleaved=False, + use_fp8=use_fp8, + sdpa_backend=sdpa_backend, + ) + self.q_proj = nn.Linear(self.query_dim, self.inner_dim, bias=False) + self.k_proj = nn.Linear(self.context_dim, self.inner_dim, bias=False) + self.v_proj = nn.Linear(self.context_dim, self.inner_dim, bias=False) + self.output_proj = nn.Linear(self.inner_dim, self.query_dim, bias=False) + self.q_norm = nn.RMSNorm(self.head_dim, eps=self.qk_norm_eps) + self.k_norm = nn.RMSNorm(self.head_dim, eps=self.qk_norm_eps) + self._initialize_derived_weights() + + def set_context_parallel_group(self, cp_group: ProcessGroup | None) -> None: + """Reject context parallelism unsupported by Triton attention. + + Args: + cp_group: Context-parallel process group; ``None`` is a no-op. + + Raises: + NotImplementedError: ``cp_group`` is not ``None``. + """ + if cp_group is not None: + raise NotImplementedError( + "The Triton attention backend does not support context parallelism" + ) + + def is_context_parallel_enabled(self) -> bool: + """Return whether context parallelism is enabled.""" + return False + + def context_parallel_size(self) -> int: + """Return the singleton context-parallel world size.""" + return 1 + + +class TritonSelfAttention(TritonMultiHeadAttention): + """Accelerated self-attention adapted to the Omnidreams contract.""" + + @property + def query_projection(self) -> nn.Linear: + """Return the canonical query projection.""" + return self.q_proj + + @property + def key_projection(self) -> nn.Linear: + """Return the canonical key projection.""" + return self.k_proj + + @property + def value_projection(self) -> nn.Linear: + """Return the canonical value projection.""" + return self.v_proj + + @property + def output_projection(self) -> nn.Linear: + """Return the canonical output projection.""" + return self.output_proj + + @property + def query_norm(self) -> nn.Module: + """Return the canonical query normalization.""" + return self.q_norm + + @property + def key_norm(self) -> nn.Module: + """Return the canonical key normalization.""" + return self.k_norm + + def __init__( + self, + query_dim: int, + context_dim: int | None = None, + n_heads: int = 8, + head_dim: int = 64, + cp_method: Literal["ring", "ulysses"] = "ring", + sdpa_backend: SDPABackend = SDPABackend.TRITON, + qkv_fusion_option: QKVFusionOption = QKVFusionOption.FULL, + use_fp8: bool = True, + ) -> None: + """Initialize bias-free Triton self-attention. + + Args: + query_dim: Feature dimension of input and output tokens. + context_dim: Self-attention context dimension. ``None`` uses + ``query_dim``. + n_heads: Number of attention heads. + head_dim: Per-head feature dimension. + cp_method: Ignored context-parallel method retained for constructor + compatibility with Omnidreams attention. + sdpa_backend: Scaled-dot-product attention implementation. + qkv_fusion_option: Projection fusion policy. + use_fp8: Whether projection GEMMs and supported attention storage + use FP8. + + Raises: + ValueError: ``context_dim`` differs from ``query_dim``. + """ + del cp_method + context_dim = query_dim if context_dim is None else context_dim + if context_dim != query_dim: + raise ValueError( + "Triton self-attention requires context_dim to equal query_dim; " + f"got {context_dim} and {query_dim}" + ) + super().__init__( + query_dim=query_dim, + n_heads=n_heads, + head_dim=head_dim, + attention_type=AttentionType.SELF_ATTENTION, + qkv_fusion_option=qkv_fusion_option, + qk_norm_eps=1e-6, + qk_norm_scope=QKNormScope.HEAD, + rope_interleaved=False, + use_fp8=use_fp8, + sdpa_backend=sdpa_backend, + ) + self.q_proj = nn.Linear(self.query_dim, self.inner_dim, bias=False) + self.k_proj = nn.Linear(self.context_dim, self.inner_dim, bias=False) + self.v_proj = nn.Linear(self.context_dim, self.inner_dim, bias=False) + self.output_proj = nn.Linear(self.inner_dim, self.query_dim, bias=False) + self.q_norm = nn.RMSNorm(self.head_dim, eps=self.qk_norm_eps) + self.k_norm = nn.RMSNorm(self.head_dim, eps=self.qk_norm_eps) + self._initialize_derived_weights() + + def set_context_parallel_group(self, cp_group: ProcessGroup | None) -> None: + """Reject context parallelism unsupported by Triton attention. + + Args: + cp_group: Context-parallel process group; ``None`` is a no-op. + + Raises: + NotImplementedError: ``cp_group`` is not ``None``. + """ + if cp_group is not None: + raise NotImplementedError( + "The Triton attention backend does not support context parallelism" + ) + + def is_context_parallel_enabled(self) -> bool: + """Return whether context parallelism is enabled.""" + return False + + def context_parallel_size(self) -> int: + """Return the singleton context-parallel world size.""" + return 1 + + @dataclass class BlockCache: """Per-block cache container for self-attention and cross-attention.""" @@ -487,34 +714,75 @@ def __init__( adaln_lora_dim: int = 256, enable_cross_view_attn: bool = False, cp_method: Literal["ring", "ulysses"] = "ring", + self_attention_backend: AttentionBackend = AttentionBackend.OMNIDREAMS, + cross_attention_backend: AttentionBackend = AttentionBackend.OMNIDREAMS, + sdpa_backend: SDPABackend = SDPABackend.TRITON, + cross_attn_sdpa_backend: SDPABackend = SDPABackend.TRITON, + self_attn_qkv_fusion_option: QKVFusionOption = QKVFusionOption.FULL, + cross_attn_qkv_fusion_option: QKVFusionOption = QKVFusionOption.FUSE_KV, + use_fp8: bool = True, ) -> None: super().__init__() self.x_dim = x_dim self.enable_cross_view_attn = enable_cross_view_attn + self.self_attention_backend = AttentionBackend(self_attention_backend) + self.cross_attention_backend = AttentionBackend(cross_attention_backend) + self.sdpa_backend = SDPABackend(sdpa_backend) + self.cross_attn_sdpa_backend = SDPABackend(cross_attn_sdpa_backend) + self.self_attn_qkv_fusion_option = QKVFusionOption(self_attn_qkv_fusion_option) + self.cross_attn_qkv_fusion_option = QKVFusionOption( + cross_attn_qkv_fusion_option + ) + self.use_fp8 = use_fp8 # Self-attention self.layer_norm_self_attn = nn.LayerNorm( x_dim, elementwise_affine=False, eps=1e-6 ) - self.self_attn = SelfAttention( - query_dim=x_dim, - context_dim=None, - n_heads=num_heads, - head_dim=x_dim // num_heads, - cp_method=cp_method, - ) # Cross-attention self.layer_norm_cross_attn = nn.LayerNorm( x_dim, elementwise_affine=False, eps=1e-6 ) - self.cross_attn = CrossAttention( - query_dim=x_dim, - context_dim=context_dim, - n_heads=num_heads, - head_dim=x_dim // num_heads, - cp_method=cp_method, - ) + if self.self_attention_backend is AttentionBackend.OMNIDREAMS: + self.self_attn = SelfAttention( + query_dim=x_dim, + context_dim=None, + n_heads=num_heads, + head_dim=x_dim // num_heads, + cp_method=cp_method, + ) + else: + self.self_attn = TritonSelfAttention( + query_dim=x_dim, + context_dim=None, + n_heads=num_heads, + head_dim=x_dim // num_heads, + cp_method=cp_method, + sdpa_backend=self.sdpa_backend, + qkv_fusion_option=self.self_attn_qkv_fusion_option, + use_fp8=self.use_fp8, + ) + + if self.cross_attention_backend is AttentionBackend.OMNIDREAMS: + self.cross_attn = CrossAttention( + query_dim=x_dim, + context_dim=context_dim, + n_heads=num_heads, + head_dim=x_dim // num_heads, + cp_method=cp_method, + ) + else: + self.cross_attn = TritonCrossAttention( + query_dim=x_dim, + context_dim=context_dim, + n_heads=num_heads, + head_dim=x_dim // num_heads, + cp_method=cp_method, + sdpa_backend=self.cross_attn_sdpa_backend, + qkv_fusion_option=self.cross_attn_qkv_fusion_option, + use_fp8=self.use_fp8, + ) # MLP self.layer_norm_mlp = nn.LayerNorm(x_dim, elementwise_affine=False, eps=1e-6) @@ -555,13 +823,25 @@ def __init__( x_dim, elementwise_affine=True, eps=1e-6 ) # dense cross view attention - self.cross_view_attn = CrossAttention( - query_dim=x_dim, - context_dim=x_dim, - n_heads=num_heads, - head_dim=x_dim // num_heads, - cp_method=cp_method, - ) + if self.cross_attention_backend is AttentionBackend.OMNIDREAMS: + self.cross_view_attn = CrossAttention( + query_dim=x_dim, + context_dim=x_dim, + n_heads=num_heads, + head_dim=x_dim // num_heads, + cp_method=cp_method, + ) + else: + self.cross_view_attn = TritonCrossAttention( + query_dim=x_dim, + context_dim=x_dim, + n_heads=num_heads, + head_dim=x_dim // num_heads, + cp_method=cp_method, + sdpa_backend=self.cross_attn_sdpa_backend, + qkv_fusion_option=self.cross_attn_qkv_fusion_option, + use_fp8=self.use_fp8, + ) def set_context_parallel_group( self, @@ -598,7 +878,7 @@ def initialize_cache( num_views = context.shape[1] self_attn_batch_size = batch_size * num_views return BlockCache( - self_attn=self.self_attn.initialize_cache( + self_attn=self.self_attn.allocate_kv_cache( self_attn_batch_size, chunk_size, window_size, @@ -606,9 +886,10 @@ def initialize_cache( device=device, dtype=dtype, ), - cross_attn=self.cross_attn.initialize_cache(context), + cross_attn=self.cross_attn.compute_kv(context), ) + @nvtx.annotate("omnidreams.dit.block") def forward( self, x: Tensor, @@ -645,106 +926,111 @@ def forward( emb = emb.reshape(B, 1, 1, D) # Compute AdaLN modulation - if self.use_adaln_lora: - assert adaln_lora is not None, ( - "adaln_lora is required when use_adaln_lora is True" - ) - adaln_lora = adaln_lora.reshape(B, 1, 1, 3 * D) - shift_self, scale_self, gate_self = ( - self.adaln_modulation_self_attn(emb) + adaln_lora - ).chunk(3, dim=-1) - shift_cross, scale_cross, gate_cross = ( - self.adaln_modulation_cross_attn(emb) + adaln_lora - ).chunk(3, dim=-1) - shift_mlp, scale_mlp, gate_mlp = ( - self.adaln_modulation_mlp(emb) + adaln_lora - ).chunk(3, dim=-1) - else: - shift_self, scale_self, gate_self = self.adaln_modulation_self_attn( - emb - ).chunk(3, dim=-1) - shift_cross, scale_cross, gate_cross = self.adaln_modulation_cross_attn( - emb - ).chunk(3, dim=-1) - shift_mlp, scale_mlp, gate_mlp = self.adaln_modulation_mlp(emb).chunk( - 3, dim=-1 - ) - - if self.enable_cross_view_attn: - assert view_embedding_proj is not None - ( - view_shift_self, - view_scale_self, - view_gate_self, - view_shift_cross, - view_scale_cross, - view_gate_cross, - view_shift_mlp, - view_scale_mlp, - view_gate_mlp, - ) = view_embedding_proj.chunk(9, dim=-1) - - def expand_view_mod(v_mod: Tensor) -> Tensor: - return v_mod.reshape(B, V, 1, D) - - shift_self = shift_self + expand_view_mod(view_shift_self) - scale_self = scale_self + expand_view_mod(view_scale_self) - gate_self = gate_self + expand_view_mod(view_gate_self) - - shift_cross = shift_cross + expand_view_mod(view_shift_cross) - scale_cross = scale_cross + expand_view_mod(view_scale_cross) - gate_cross = gate_cross + expand_view_mod(view_gate_cross) - - shift_mlp = shift_mlp + expand_view_mod(view_shift_mlp) - scale_mlp = scale_mlp + expand_view_mod(view_scale_mlp) - gate_mlp = gate_mlp + expand_view_mod(view_gate_mlp) + with nvtx.annotate("omnidreams.dit.adaln"): + if self.use_adaln_lora: + assert adaln_lora is not None, ( + "adaln_lora is required when use_adaln_lora is True" + ) + adaln_lora = adaln_lora.reshape(B, 1, 1, 3 * D) + shift_self, scale_self, gate_self = ( + self.adaln_modulation_self_attn(emb) + adaln_lora + ).chunk(3, dim=-1) + shift_cross, scale_cross, gate_cross = ( + self.adaln_modulation_cross_attn(emb) + adaln_lora + ).chunk(3, dim=-1) + shift_mlp, scale_mlp, gate_mlp = ( + self.adaln_modulation_mlp(emb) + adaln_lora + ).chunk(3, dim=-1) + else: + shift_self, scale_self, gate_self = self.adaln_modulation_self_attn( + emb + ).chunk(3, dim=-1) + shift_cross, scale_cross, gate_cross = self.adaln_modulation_cross_attn( + emb + ).chunk(3, dim=-1) + shift_mlp, scale_mlp, gate_mlp = self.adaln_modulation_mlp(emb).chunk( + 3, dim=-1 + ) + + if self.enable_cross_view_attn: + assert view_embedding_proj is not None + ( + view_shift_self, + view_scale_self, + view_gate_self, + view_shift_cross, + view_scale_cross, + view_gate_cross, + view_shift_mlp, + view_scale_mlp, + view_gate_mlp, + ) = view_embedding_proj.chunk(9, dim=-1) + + def expand_view_mod(v_mod: Tensor) -> Tensor: + return v_mod.reshape(B, V, 1, D) + + shift_self = shift_self + expand_view_mod(view_shift_self) + scale_self = scale_self + expand_view_mod(view_scale_self) + gate_self = gate_self + expand_view_mod(view_gate_self) + + shift_cross = shift_cross + expand_view_mod(view_shift_cross) + scale_cross = scale_cross + expand_view_mod(view_scale_cross) + gate_cross = gate_cross + expand_view_mod(view_gate_cross) + + shift_mlp = shift_mlp + expand_view_mod(view_shift_mlp) + scale_mlp = scale_mlp + expand_view_mod(view_scale_mlp) + gate_mlp = gate_mlp + expand_view_mod(view_gate_mlp) # Self-attention - normed_x = self.layer_norm_self_attn(x) * (1 + scale_self) + shift_self - attn_out = self.self_attn( - normed_x, - rope_freqs=rope_freqs, - kv_cache=cache.self_attn, - ).reshape_as(normed_x) - x = x + gate_self * attn_out + with nvtx.annotate("omnidreams.dit.self_attention"): + normed_x = self.layer_norm_self_attn(x) * (1 + scale_self) + shift_self + attn_out = self.self_attn( + normed_x, + rope_freqs=rope_freqs, + kv_cache=cache.self_attn, + ).reshape_as(normed_x) + x = x + gate_self * attn_out # Cross-view attention: dense if self.enable_cross_view_attn: - assert T is not None and HW is not None, ( - "T and HW must be available (x should be a 5D tensor) when cross-view attention is enabled" - ) - normed_x_cv = self.layer_norm_cross_view_attn(x) - x_cv = rearrange(normed_x_cv, "b v (t hw) d -> b t v hw d", t=T, hw=HW) - if self.cross_view_attn.is_context_parallel_enabled(): - # CP-enabled: views are split across GPUs in rank order - # (e.g. 4 views on 2 GPUs -> [0,1] and [2,3]). - if V == 1: - # CP size == num views: ring attention gathers all K/V, - # so local context stays unexpanded. - x_context = x_cv + with nvtx.annotate("omnidreams.dit.cross_view_attention"): + assert T is not None and HW is not None, ( + "T and HW must be available (x should be a 5D tensor) when cross-view attention is enabled" + ) + normed_x_cv = self.layer_norm_cross_view_attn(x) + x_cv = rearrange(normed_x_cv, "b v (t hw) d -> b t v hw d", t=T, hw=HW) + if self.cross_view_attn.is_context_parallel_enabled(): + # CP-enabled: views are split across GPUs in rank order + # (e.g. 4 views on 2 GPUs -> [0,1] and [2,3]). + if V == 1: + # CP size == num views: ring attention gathers all K/V, + # so local context stays unexpanded. + x_context = x_cv + else: + # CP size < num views: gather each GPU's local views first. + x_context = repeat(x_cv, "b t v hw d -> b t v2 (v hw) d", v2=V) else: - # CP size < num views: gather each GPU's local views first. + # Without CP, repeat context so each view attends over all views. x_context = repeat(x_cv, "b t v hw d -> b t v2 (v hw) d", v2=V) - else: - # Without CP, repeat context so each view attends over all views. - x_context = repeat(x_cv, "b t v hw d -> b t v2 (v hw) d", v2=V) - cross_view_attn_kv_cache = self.cross_view_attn.compute_kv(x_context) - cv_out = self.cross_view_attn(x_cv, kv_cache=cross_view_attn_kv_cache) - cv_out = rearrange(cv_out, "b t v hw d -> b v (t hw) d") - x = x + cv_out + cross_view_attn_kv_cache = self.cross_view_attn.compute_kv(x_context) + cv_out = self.cross_view_attn(x_cv, kv_cache=cross_view_attn_kv_cache) + cv_out = rearrange(cv_out, "b t v hw d -> b v (t hw) d") + x = x + cv_out # Cross-attention - normed_x = self.layer_norm_cross_attn(x) * (1 + scale_cross) + shift_cross - cross_out = self.cross_attn( - normed_x, - kv_cache=cache.cross_attn, - ).reshape_as(normed_x) - x = x + gate_cross * cross_out + with nvtx.annotate("omnidreams.dit.cross_attention"): + normed_x = self.layer_norm_cross_attn(x) * (1 + scale_cross) + shift_cross + cross_out = self.cross_attn( + normed_x, + kv_cache=cache.cross_attn, + ).reshape_as(normed_x) + x = x + gate_cross * cross_out # MLP - normed_x = self.layer_norm_mlp(x) * (1 + scale_mlp) + shift_mlp - mlp_out = self.mlp(normed_x) - x = x + gate_mlp * mlp_out + with nvtx.annotate("omnidreams.dit.mlp"): + normed_x = self.layer_norm_mlp(x) * (1 + scale_mlp) + shift_mlp + mlp_out = self.mlp(normed_x) + x = x + gate_mlp * mlp_out # reshape back to 5D if needed if T is not None and HW is not None: diff --git a/integrations/omnidreams/omnidreams/transformer/impl/network.py b/integrations/omnidreams/omnidreams/transformer/impl/network.py index 307e88220..96b6a4a4d 100644 --- a/integrations/omnidreams/omnidreams/transformer/impl/network.py +++ b/integrations/omnidreams/omnidreams/transformer/impl/network.py @@ -24,6 +24,10 @@ from torch import Tensor from torch.distributed import ProcessGroup +from flashdreams.accelerated.multi_head_attention_triton import ( + QKVFusionOption, + SDPABackend, +) from flashdreams.core.distributed.context_parallel import ( cat_outputs_cp, split_inputs_cp, @@ -31,6 +35,7 @@ from flashdreams.infra.config import InstantiateConfig from .modules import ( + AttentionBackend, Block, BlockCache, FinalLayer, @@ -119,6 +124,27 @@ class CosmosDiTNetworkConfig(InstantiateConfig): cp_method: Literal["ring", "ulysses"] = "ring" """Context-parallel attention method for transformer attention ops.""" + self_attention_backend: AttentionBackend = AttentionBackend.OMNIDREAMS + """Self-attention implementation used by every DiT block.""" + + cross_attention_backend: AttentionBackend = AttentionBackend.OMNIDREAMS + """Text and cross-view attention implementation used by every DiT block.""" + + sdpa_backend: SDPABackend = SDPABackend.TRITON + """SDPA implementation used by accelerated self-attention.""" + + cross_attn_sdpa_backend: SDPABackend = SDPABackend.TRITON + """SDPA implementation used by accelerated cross-attention.""" + + self_attn_qkv_fusion_option: QKVFusionOption = QKVFusionOption.FULL + """Projection fusion policy used by accelerated self-attention.""" + + cross_attn_qkv_fusion_option: QKVFusionOption = QKVFusionOption.FUSE_KV + """Projection fusion policy used by accelerated cross-attention.""" + + use_fp8: bool = True + """Whether accelerated attention projections and supported storage use FP8.""" + view_condition_dim: int = 16 """Embedding dim for the per-view conditioning vector.""" @@ -132,6 +158,15 @@ class CosmosDiTNetwork(nn.Module): def __init__(self, config: CosmosDiTNetworkConfig): super().__init__() self.config = config + self.sdpa_backend = SDPABackend(config.sdpa_backend) + self.cross_attn_sdpa_backend = SDPABackend(config.cross_attn_sdpa_backend) + self.self_attn_qkv_fusion_option = QKVFusionOption( + config.self_attn_qkv_fusion_option + ) + self.cross_attn_qkv_fusion_option = QKVFusionOption( + config.cross_attn_qkv_fusion_option + ) + self.use_fp8 = config.use_fp8 # add 1 for the condition mask in_channels = config.in_channels + 1 @@ -177,6 +212,13 @@ def __init__(self, config: CosmosDiTNetworkConfig): adaln_lora_dim=self.config.adaln_lora_dim, enable_cross_view_attn=self.config.enable_cross_view_attn, cp_method=self.config.cp_method, + self_attention_backend=self.config.self_attention_backend, + cross_attention_backend=self.config.cross_attention_backend, + sdpa_backend=self.sdpa_backend, + cross_attn_sdpa_backend=self.cross_attn_sdpa_backend, + self_attn_qkv_fusion_option=self.self_attn_qkv_fusion_option, + cross_attn_qkv_fusion_option=self.cross_attn_qkv_fusion_option, + use_fp8=self.use_fp8, ) for _ in range(self.config.num_blocks) ] diff --git a/integrations/omnidreams/pyproject.toml b/integrations/omnidreams/pyproject.toml index 2e38c80fb..52175dd85 100644 --- a/integrations/omnidreams/pyproject.toml +++ b/integrations/omnidreams/pyproject.toml @@ -119,6 +119,9 @@ interactive-drive-configuration = "omnidreams.interactive_drive.input_config.app [project.entry-points."flashdreams.runner_configs"] "omnidreams" = "omnidreams.config:RUNNER_SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE" "omnidreams-perf" = "omnidreams.config:RUNNER_SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE_PERF" +"omnidreams-cuda-cudnn" = "omnidreams.config:RUNNER_SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE_CUDA_CUDNN" +"omnidreams-cuda-sparge" = "omnidreams.config:RUNNER_SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE_CUDA_SPARGE" +"omnidreams-cuda-sage3fp8" = "omnidreams.config:RUNNER_SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE_CUDA_SAGE3_FP8" [tool.setuptools.packages.find] include = ["omnidreams*"] diff --git a/integrations/omnidreams/tests/interactive_drive/test_presenter.py b/integrations/omnidreams/tests/interactive_drive/test_presenter.py index b3c23b2ec..3a9d4a2f5 100644 --- a/integrations/omnidreams/tests/interactive_drive/test_presenter.py +++ b/integrations/omnidreams/tests/interactive_drive/test_presenter.py @@ -989,6 +989,27 @@ def test_hud_resize_uses_actual_window_size_without_model_resolution_clamp() -> def test_hud_auto_sizes_window_to_native_model_frame_resolution() -> None: presenter = _hud_presenter_without_window() resize_calls: list[tuple[int, int]] = [] + presenter._native_model_auto_resize_enabled = True + presenter._auto_sized_camera_src_size = None + presenter._pending_resize = None + presenter._window = SimpleNamespace( + size=SimpleNamespace(x=1920, y=1080), + resize=lambda width, height: resize_calls.append((width, height)), + ) + + resized = presenter._resize_window_for_native_model_frame( + np.zeros((1200, 1600, 3), dtype=np.uint8) + ) + + assert resized is True + assert resize_calls == [(2100, 1200)] + assert presenter._pending_resize == (2100, 1200) + + +def test_hud_does_not_grow_window_clamped_during_initialization() -> None: + presenter = _hud_presenter_without_window() + resize_calls: list[tuple[int, int]] = [] + presenter._native_model_auto_resize_enabled = True presenter._auto_sized_camera_src_size = None presenter._pending_resize = None presenter._window = SimpleNamespace( @@ -1000,9 +1021,10 @@ def test_hud_auto_sizes_window_to_native_model_frame_resolution() -> None: np.zeros((704, 1280, 3), dtype=np.uint8) ) - assert resized is True - assert resize_calls == [(1780, 704)] - assert presenter._pending_resize == (1780, 704) + assert resized is False + assert resize_calls == [] + assert presenter._pending_resize is None + assert presenter._auto_sized_camera_src_size == (1280, 704) def test_hud_keeps_larger_canvas_when_model_resolution_shrinks() -> None: @@ -1182,6 +1204,60 @@ def raise_cuda(frame: PresentedFrame, rgb: object) -> bool: assert close_calls == 1 +def test_hud_close_releases_slangpy_resources_in_dependency_order() -> None: + presenter = _hud_presenter_without_window() + events: list[str] = [] + + class _Interop: + def close(self) -> None: + events.append("interop.close") + + class _Device: + def wait_for_idle(self) -> None: + events.append("device.wait_for_idle") + + def close(self) -> None: + events.append("device.close") + + class _Surface: + def unconfigure(self) -> None: + events.append("surface.unconfigure") + + class _Window: + def close(self) -> None: + events.append("window.close") + + presenter._bev_panel_exec = None + presenter._cuda_hud_interop = _Interop() + presenter._retired_cuda_hud_interops = [_Interop()] + presenter._wheel = None + presenter._device = _Device() + presenter._surface = _Surface() + presenter._camera_texture = object() + presenter._camera_fit_texture = object() + presenter._display_texture = object() + presenter._window = _Window() + + presenter.close() + presenter.close() + + assert events == [ + "interop.close", + "interop.close", + "device.wait_for_idle", + "surface.unconfigure", + "device.close", + "window.close", + ] + assert presenter._retired_cuda_hud_interops == [] + assert presenter._camera_texture is None + assert presenter._camera_fit_texture is None + assert presenter._display_texture is None + assert presenter._surface is None + assert presenter._device is None + assert presenter._window is None + + class _ExitSceneKeyboard: def __init__(self) -> None: self.cleared = 0 diff --git a/integrations/omnidreams/tests/test_omnidreams_singleview_native.py b/integrations/omnidreams/tests/test_omnidreams_singleview_native.py index 66193fc70..a9ed6ab43 100644 --- a/integrations/omnidreams/tests/test_omnidreams_singleview_native.py +++ b/integrations/omnidreams/tests/test_omnidreams_singleview_native.py @@ -153,6 +153,7 @@ def fake_load_torch_extension(**kwargs: object) -> ModuleType: monkeypatch.setattr(cpp_extension, "load", fake_load_torch_extension) monkeypatch.setattr(native.os, "cpu_count", lambda: 48) monkeypatch.setattr(native, "_python_package_dir", lambda package: None) + monkeypatch.setattr(native, "_detected_cuda_arch_list", lambda: None) monkeypatch.delenv("MAX_JOBS", raising=False) monkeypatch.delenv("TORCH_CUDA_ARCH_LIST", raising=False) monkeypatch.delenv("OMNIDREAMS_SINGLEVIEW_CUDA_ARCH_LIST", raising=False) @@ -211,8 +212,6 @@ def fake_load_torch_extension(**kwargs: object) -> ModuleType: "lightvae_fp8_warp_mma_stages.cu", "lightvae_fp8_attention.cu", "streaming_dit_bridge.cu", - "sage3_blackwell_api_shim.cu", - "sage3_fp4_quant_shim.cu", "attention.cu", "block_quant.cu", "cosmos_adaln_lora.cu", @@ -224,7 +223,7 @@ def fake_load_torch_extension(**kwargs: object) -> ModuleType: "cosmos_gemm_bf16.cu", "cosmos_modulate.cu", "ops.cu", - "sage3_attention.cu", + "sage3_attention_stub.cu", "sparge_attention_sm89_inst.cu", "transformer_block.cu", ] @@ -267,27 +266,81 @@ def fake_load_torch_extension(**kwargs: object) -> ModuleType: '-DOMNIDREAMS_SINGLEVIEW_SAGE_ATTENTION_SHA=\\"sage-test-sha\\"' in captured["extra_cflags"] ) - assert "-DOMNIDREAMS_SINGLEVIEW_HAS_SAGE3=1" in captured["extra_cflags"] + assert "-DOMNIDREAMS_SINGLEVIEW_HAS_SAGE3=0" in captured["extra_cflags"] assert ( '-DOMNIDREAMS_SINGLEVIEW_SPARGE_ATTN_SHA=\\"sparge-test-sha\\"' in captured["extra_cflags"] ) assert "-DOMNIDREAMS_SINGLEVIEW_HAS_SPARGE=1" in captured["extra_cflags"] assert ( - '-DOMNIDREAMS_SINGLEVIEW_CUDA_ARCH_LIST=\\"12.0a\\"' in captured["extra_cflags"] + '-DOMNIDREAMS_SINGLEVIEW_CUDA_ARCH_LIST=\\"pytorch-default\\"' + in captured["extra_cflags"] ) assert "-DOMNIDREAMS_SINGLEVIEW_WITH_CUDA" in captured["extra_cuda_cflags"] if os.name == "nt": assert "-Xcompiler=/Zc:preprocessor" in captured["extra_cuda_cflags"] - assert "-DOMNIDREAMS_SINGLEVIEW_HAS_SAGE3=1" in captured["extra_cuda_cflags"] + assert "-DOMNIDREAMS_SINGLEVIEW_HAS_SAGE3=0" in captured["extra_cuda_cflags"] assert "-DOMNIDREAMS_SINGLEVIEW_HAS_SPARGE=1" in captured["extra_cuda_cflags"] assert captured["with_cuda"] is True assert captured["max_jobs_env"] == "8" - assert captured["cuda_arch_list_env"] == "12.0a" + assert captured["cuda_arch_list_env"] is None assert "MAX_JOBS" not in os.environ assert "TORCH_CUDA_ARCH_LIST" not in os.environ +@pytest.mark.ci_cpu +@pytest.mark.parametrize( + ("capability", "device_name", "expected"), + [ + ((12, 0), "NVIDIA GeForce RTX 5090", "12.0a"), + ((12, 0), "NVIDIA RTX PRO 6000 Blackwell", "12.0a"), + ((12, 0), "Unvalidated Compute Capability 12.0 GPU", None), + ((10, 3), "NVIDIA GB300", None), + ((8, 9), "NVIDIA RTX 6000 Ada Generation", None), + ], +) +def test_detected_cuda_arch_list_only_selects_validated_sm120a_devices( + capability: tuple[int, int], + device_name: str, + expected: str | None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda: capability) + monkeypatch.setattr(torch.cuda, "get_device_name", lambda: device_name) + + assert native._detected_cuda_arch_list() == expected + + +@pytest.mark.ci_cpu +def test_effective_cuda_arch_list_precedence( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(native, "_detected_cuda_arch_list", lambda: "12.0a") + monkeypatch.delenv("TORCH_CUDA_ARCH_LIST", raising=False) + monkeypatch.delenv("OMNIDREAMS_SINGLEVIEW_CUDA_ARCH_LIST", raising=False) + + assert native._effective_cuda_arch_list() == "12.0a" + + monkeypatch.setenv("OMNIDREAMS_SINGLEVIEW_CUDA_ARCH_LIST", "10.3a") + assert native._effective_cuda_arch_list() == "10.3a" + + monkeypatch.setenv("TORCH_CUDA_ARCH_LIST", "8.9") + assert native._effective_cuda_arch_list() == "8.9" + + +@pytest.mark.ci_cpu +def test_effective_cuda_arch_list_uses_pytorch_default_without_sm120a( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(native, "_detected_cuda_arch_list", lambda: None) + monkeypatch.delenv("TORCH_CUDA_ARCH_LIST", raising=False) + monkeypatch.delenv("OMNIDREAMS_SINGLEVIEW_CUDA_ARCH_LIST", raising=False) + + assert native._effective_cuda_arch_list() is None + assert native._cuda_arch_identity(None) == "pytorch-default" + + @pytest.mark.ci_cpu @pytest.mark.parametrize( ("value", "expected"), @@ -307,6 +360,7 @@ def test_sage3_build_opt_out_parses_affirmative_values( expected: bool, monkeypatch: pytest.MonkeyPatch, ) -> None: + monkeypatch.setenv("OMNIDREAMS_SINGLEVIEW_CUDA_ARCH_LIST", "12.0a") if value is None: monkeypatch.delenv("OMNIDREAMS_SINGLEVIEW_DISABLE_SAGE3", raising=False) else: @@ -326,6 +380,23 @@ def test_sage3_build_opt_out_parses_affirmative_values( assert "sage3_attention.cu" in sources +@pytest.mark.ci_cpu +@pytest.mark.parametrize( + ("cuda_arch_list", "expected"), + [ + ("12.0a", False), + ("pytorch-default", True), + ("10.3a", True), + ("12.0", True), + ], +) +def test_sage3_build_requires_exact_sm120a_target( + cuda_arch_list: str, + expected: bool, +) -> None: + assert native._sage3_disabled(cuda_arch_list) is expected + + @pytest.mark.ci_cpu def test_load_extension_uses_sage3_stub_when_disabled( tmp_path: Path, @@ -381,6 +452,7 @@ def fake_load_torch_extension(**_: object) -> ModuleType: monkeypatch.setattr(native, "validate_thirdparty", lambda: thirdparty_info) monkeypatch.setattr(cpp_extension, "load", fake_load_torch_extension) monkeypatch.setattr(native, "_python_package_dir", lambda package: None) + monkeypatch.setenv("OMNIDREAMS_SINGLEVIEW_CUDA_ARCH_LIST", "12.0a") monkeypatch.delenv("OMNIDREAMS_SINGLEVIEW_DISABLE_SAGE3", raising=False) sage3_extension = native.load_extension(build_root=tmp_path / "native-build") @@ -397,6 +469,45 @@ def fake_load_torch_extension(**_: object) -> ModuleType: assert len(extensions) == 2 +@pytest.mark.ci_cpu +def test_load_extension_caches_separate_cuda_architectures( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + import torch.utils.cpp_extension as cpp_extension + + extensions: list[ModuleType] = [] + + def fake_load_torch_extension(**_: object) -> ModuleType: + extension = _fake_extension_module() + extensions.append(extension) + return extension + + thirdparty_info = _fake_thirdparty_info(tmp_path) + monkeypatch.setattr(native, "_extension", {}) + monkeypatch.setattr(native, "_extension_load_error", None) + monkeypatch.setattr(native, "validate_thirdparty", lambda: thirdparty_info) + monkeypatch.setattr(cpp_extension, "load", fake_load_torch_extension) + monkeypatch.setattr(native, "_python_package_dir", lambda package: None) + monkeypatch.delenv("TORCH_CUDA_ARCH_LIST", raising=False) + + monkeypatch.setenv("OMNIDREAMS_SINGLEVIEW_CUDA_ARCH_LIST", "10.3a") + sm103_extension = native.load_extension(build_root=tmp_path / "native-build") + monkeypatch.setenv("OMNIDREAMS_SINGLEVIEW_CUDA_ARCH_LIST", "12.0a") + sm120_extension = native.load_extension(build_root=tmp_path / "native-build") + + assert sm103_extension is extensions[0] + assert sm120_extension is extensions[1] + assert ( + native.load_extension(build_root=tmp_path / "native-build") is sm120_extension + ) + monkeypatch.setenv("OMNIDREAMS_SINGLEVIEW_CUDA_ARCH_LIST", "10.3a") + assert ( + native.load_extension(build_root=tmp_path / "native-build") is sm103_extension + ) + assert len(extensions) == 2 + + @pytest.mark.ci_cpu def test_extension_name_isolated_by_sage3_build_opt_out( tmp_path: Path, @@ -404,6 +515,7 @@ def test_extension_name_isolated_by_sage3_build_opt_out( ) -> None: thirdparty_info = _fake_thirdparty_info(tmp_path) monkeypatch.setattr(native, "_source_fingerprint", lambda: "fixed-fingerprint") + monkeypatch.setenv("OMNIDREAMS_SINGLEVIEW_CUDA_ARCH_LIST", "12.0a") monkeypatch.delenv("OMNIDREAMS_SINGLEVIEW_DISABLE_SAGE3", raising=False) full_name = native._extension_name(thirdparty_info) @@ -415,6 +527,26 @@ def test_extension_name_isolated_by_sage3_build_opt_out( assert "_sage3_0_" in stubbed_name +@pytest.mark.ci_cpu +def test_extension_name_isolated_by_cuda_architecture( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + thirdparty_info = _fake_thirdparty_info(tmp_path) + monkeypatch.setattr(native, "_source_fingerprint", lambda: "fixed-fingerprint") + + sm103_name = native._extension_name( + thirdparty_info, + cuda_arch_list="10.3a", + ) + sm120_name = native._extension_name( + thirdparty_info, + cuda_arch_list="12.0a", + ) + + assert sm103_name != sm120_name + + @pytest.mark.ci_cpu def test_load_extension_respects_existing_max_jobs( tmp_path: Path, @@ -937,10 +1069,7 @@ def test_cuda_native_extension_builds(tmp_path: Path) -> None: assert extension.is_available() build_info = extension.build_info() assert build_info["with_cuda"] is True - expected_arch = os.environ.get( - "TORCH_CUDA_ARCH_LIST", - os.environ.get("OMNIDREAMS_SINGLEVIEW_CUDA_ARCH_LIST", "12.0a"), - ) + expected_arch = native._cuda_arch_identity(native._effective_cuda_arch_list()) assert build_info["cuda_arch_list"] == expected_arch assert hasattr(extension, "native_tensor_descriptor") assert hasattr(extension, "native_tensor_ref_descriptor") diff --git a/integrations/omnidreams/tests/test_recipe_configs.py b/integrations/omnidreams/tests/test_recipe_configs.py index 15e4004c6..8ebbb788b 100644 --- a/integrations/omnidreams/tests/test_recipe_configs.py +++ b/integrations/omnidreams/tests/test_recipe_configs.py @@ -51,6 +51,15 @@ def test_public_runner_slugs_map_to_internal_pipeline_presets() -> None: expected = { "omnidreams": "omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae", "omnidreams-perf": "omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae-perf", + "omnidreams-cuda-cudnn": ( + "omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae-cuda-cudnn" + ), + "omnidreams-cuda-sparge": ( + "omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae-cuda-sparge" + ), + "omnidreams-cuda-sage3fp8": ( + "omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae-cuda-sage3-fp8" + ), } actual = {slug: cfg.pipeline.name for slug, cfg in OMNIDREAMS_RUNNERS.items()} assert actual == expected diff --git a/integrations/omnidreams/tests/test_transformer_attention_backend.py b/integrations/omnidreams/tests/test_transformer_attention_backend.py new file mode 100644 index 000000000..90106da86 --- /dev/null +++ b/integrations/omnidreams/tests/test_transformer_attention_backend.py @@ -0,0 +1,441 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""CPU coverage for Omnidreams DiT attention backend selection.""" + +import pytest +import torch +from omnidreams.transformer.impl import modules as transformer_modules +from omnidreams.transformer.impl.modules import AttentionBackend, Block +from omnidreams.transformer.impl.network import CosmosDiTNetwork, CosmosDiTNetworkConfig + +from flashdreams.accelerated import multi_head_attention_triton as triton_attention +from flashdreams.accelerated.multi_head_attention import ( + AttentionType, +) +from flashdreams.accelerated.multi_head_attention_triton import ( + QKVFusionOption, + SDPABackend, + TritonMultiHeadAttention, +) +from integrations.omnidreams.benchmarks.cases import BENCHMARK_CASES +from integrations.omnidreams.benchmarks.test_modules import ( + _MODULE_BENCHMARK_CASES, + _MODULE_CROSS_ATTENTION_CASES, + _MODULE_SELF_ATTENTION_CASES, +) + +pytestmark = pytest.mark.ci_cpu + + +def test_dit_attention_backend_defaults_to_omnidreams() -> None: + """Keep existing Omnidreams attention as the default.""" + default_block = Block( + x_dim=12, + context_dim=8, + num_heads=1, + enable_cross_view_attn=True, + ) + + assert default_block.self_attention_backend is AttentionBackend.OMNIDREAMS + assert default_block.cross_attention_backend is AttentionBackend.OMNIDREAMS + assert transformer_modules.MultiHeadAttention.__base__ is torch.nn.Module + assert isinstance(default_block.self_attn, transformer_modules.SelfAttention) + assert isinstance(default_block.cross_attn, transformer_modules.CrossAttention) + + +@pytest.mark.parametrize( + ("self_attention_backend", "cross_attention_backend"), + ( + (AttentionBackend.TRITON, AttentionBackend.OMNIDREAMS), + (AttentionBackend.OMNIDREAMS, AttentionBackend.TRITON), + ), + ids=("triton-self", "triton-cross"), +) +def test_network_config_selects_attention_backends_independently( + self_attention_backend: AttentionBackend, + cross_attention_backend: AttentionBackend, +) -> None: + """Select self- and cross-attention implementations independently.""" + config = CosmosDiTNetworkConfig( + model_channels=32, + num_blocks=1, + num_heads=2, + crossattn_emb_channels=16, + use_crossattn_projection=False, + enable_cross_view_attn=True, + self_attention_backend=self_attention_backend, + cross_attention_backend=cross_attention_backend, + sdpa_backend=SDPABackend.CUDNN, + use_fp8=False, + ) + + block = CosmosDiTNetwork(config).blocks[0] + + expected_self_type = ( + transformer_modules.TritonSelfAttention + if self_attention_backend is AttentionBackend.TRITON + else transformer_modules.SelfAttention + ) + expected_cross_type = ( + transformer_modules.TritonCrossAttention + if cross_attention_backend is AttentionBackend.TRITON + else transformer_modules.CrossAttention + ) + assert block.self_attention_backend is self_attention_backend + assert block.cross_attention_backend is cross_attention_backend + assert isinstance(block.self_attn, expected_self_type) + assert isinstance(block.cross_attn, expected_cross_type) + assert isinstance(block.cross_view_attn, expected_cross_type) + + +def test_omnidreams_attention_preserves_cache_lifecycles( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Update self-attention cache while keeping cross-attention cache static.""" + + def cpu_sdpa( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + ) -> torch.Tensor: + return torch.nn.functional.scaled_dot_product_attention(query, key, value) + + self_attention = transformer_modules.SelfAttention( + query_dim=16, + n_heads=1, + head_dim=16, + ) + monkeypatch.setattr( + transformer_modules, "apply_rope_freqs", lambda tensor, _: tensor + ) + monkeypatch.setattr(self_attention.attn_op, "_impl", cpu_sdpa) + self_cache = self_attention.allocate_kv_cache( + batch_size=2, + chunk_size=3, + window_size=6, + sink_size=0, + device=torch.device("cpu"), + dtype=torch.float32, + ) + query = torch.randn(1, 2, 3, 16) + self_cache.before_update(0) + output = self_attention( + query, + self_cache, + rope_freqs=torch.zeros(3, 1, 1, 16), + ) + assert self_cache.cached_k().shape == (2, 3, 1, 16) + self_cache.after_update(0) + assert output.shape == query.shape + + cross_attention = transformer_modules.CrossAttention( + query_dim=16, + context_dim=16, + n_heads=1, + head_dim=16, + ) + monkeypatch.setattr(cross_attention.attn_op, "_impl", cpu_sdpa) + cross_cache = cross_attention.compute_kv(torch.randn(1, 2, 5, 16)) + cached_key = cross_cache.cached_k().clone() + cached_value = cross_cache.cached_v().clone() + output = cross_attention(query, cross_cache) + assert output.shape == query.shape + torch.testing.assert_close(cross_cache.cached_k(), cached_key) + torch.testing.assert_close(cross_cache.cached_v(), cached_value) + + +@pytest.mark.parametrize( + "sdpa_backend", tuple(SDPABackend), ids=lambda backend: backend.value +) +def test_network_config_selects_triton_attention( + sdpa_backend: SDPABackend, +) -> None: + """Propagate the configured self-attention SDPA implementation.""" + config = CosmosDiTNetworkConfig( + model_channels=32, + num_blocks=1, + num_heads=2, + crossattn_emb_channels=16, + use_crossattn_projection=False, + enable_cross_view_attn=True, + self_attention_backend=AttentionBackend.TRITON, + cross_attention_backend=AttentionBackend.TRITON, + sdpa_backend=sdpa_backend, + ) + + network = CosmosDiTNetwork(config) + block = network.blocks[0] + + assert config.sdpa_backend is sdpa_backend + assert network.sdpa_backend is sdpa_backend + assert block.sdpa_backend is sdpa_backend + self_attention = block.self_attn + assert isinstance(self_attention, TritonMultiHeadAttention) + assert self_attention.use_fp8 is True + assert self_attention.attention_type is AttentionType.SELF_ATTENTION + assert self_attention.qkv_fusion_option is QKVFusionOption.FULL + assert self_attention.sdpa_backend is sdpa_backend + assert self_attention._derived_weights.fused_qkv_weight is not None + assert self_attention._derived_weights.fused_qkv_weight.dtype == torch.float8_e4m3fn + assert self_attention._derived_weights.output_weight_fp8 is not None + assert ( + self_attention._derived_weights.output_weight_fp8.dtype == torch.float8_e4m3fn + ) + cache = self_attention.allocate_kv_cache( + batch_size=1, + chunk_size=2, + window_size=4, + sink_size=0, + device=torch.device("cpu"), + dtype=torch.bfloat16, + ) + assert cache.dtype is ( + torch.float8_e4m3fn if sdpa_backend is SDPABackend.TRITON else torch.bfloat16 + ) + with pytest.raises(TypeError, match="FP8 projections require"): + self_attention.allocate_kv_cache( + batch_size=1, + chunk_size=2, + window_size=4, + sink_size=0, + device=torch.device("cpu"), + dtype=torch.float32, + ) + assert isinstance(block.cross_attn, transformer_modules.TritonCrossAttention) + assert isinstance(block.cross_view_attn, transformer_modules.TritonCrossAttention) + assert block.cross_attn.context_dim == 16 + assert block.cross_attn.attention_type is AttentionType.CROSS_ATTENTION + assert block.cross_attn.qkv_fusion_option is QKVFusionOption.FUSE_KV + assert block.cross_attn.use_fp8 is True + assert block.cross_attn.sdpa_backend is SDPABackend.TRITON + assert block.cross_view_attn.context_dim == 32 + assert block.cross_view_attn.attention_type is AttentionType.CROSS_ATTENTION + assert block.cross_view_attn.qkv_fusion_option is QKVFusionOption.FUSE_KV + assert block.cross_view_attn.use_fp8 is True + assert block.cross_view_attn.sdpa_backend is SDPABackend.TRITON + + +def test_network_config_selects_triton_attention_policies() -> None: + """Propagate cross-attention SDPA, QKV fusion, and FP8 policies.""" + config = CosmosDiTNetworkConfig( + model_channels=32, + num_blocks=1, + num_heads=2, + crossattn_emb_channels=16, + use_crossattn_projection=False, + enable_cross_view_attn=True, + self_attention_backend=AttentionBackend.TRITON, + cross_attention_backend=AttentionBackend.TRITON, + cross_attn_sdpa_backend=SDPABackend.CUDNN, + self_attn_qkv_fusion_option=QKVFusionOption.NONE, + cross_attn_qkv_fusion_option=QKVFusionOption.NONE, + use_fp8=False, + ) + + network = CosmosDiTNetwork(config) + block = network.blocks[0] + assert isinstance(block, Block) + self_attention = block.self_attn + cross_attention = block.cross_attn + cross_view_attention = block.cross_view_attn + assert isinstance(self_attention, TritonMultiHeadAttention) + assert isinstance(cross_attention, TritonMultiHeadAttention) + assert isinstance(cross_view_attention, TritonMultiHeadAttention) + + assert network.cross_attn_sdpa_backend is SDPABackend.CUDNN + assert network.self_attn_qkv_fusion_option is QKVFusionOption.NONE + assert network.cross_attn_qkv_fusion_option is QKVFusionOption.NONE + assert network.use_fp8 is False + assert self_attention.sdpa_backend is SDPABackend.TRITON + assert cross_attention.sdpa_backend is SDPABackend.CUDNN + assert cross_view_attention.sdpa_backend is SDPABackend.CUDNN + assert self_attention.qkv_fusion_option is QKVFusionOption.NONE + assert cross_attention.qkv_fusion_option is QKVFusionOption.NONE + assert cross_view_attention.qkv_fusion_option is QKVFusionOption.NONE + assert self_attention.use_fp8 is False + assert cross_attention.use_fp8 is False + assert cross_view_attention.use_fp8 is False + + +def test_benchmark_cases_match_selected_matrix() -> None: + """Keep the end-to-end benchmark matrix limited to selected configurations.""" + pytorch_cases = [case for case in BENCHMARK_CASES if not case.native_dit] + assert tuple( + ( + case.implementation, + case.self_attention_backend, + case.cross_attention_backend, + case.sdpa_backend, + case.use_fp8, + case.self_attn_qkv_fusion_option, + case.cross_attn_qkv_fusion_option, + ) + for case in pytorch_cases + ) == ( + ( + "omnidreams_torch", + AttentionBackend.OMNIDREAMS, + AttentionBackend.OMNIDREAMS, + SDPABackend.CUDNN, + False, + QKVFusionOption.NONE, + QKVFusionOption.NONE, + ), + ( + "triton_fa2_fp8_full", + AttentionBackend.TRITON, + AttentionBackend.TRITON, + SDPABackend.TRITON, + True, + QKVFusionOption.FULL, + QKVFusionOption.FUSE_KV, + ), + ( + "triton_cudnn_bf16_full", + AttentionBackend.TRITON, + AttentionBackend.TRITON, + SDPABackend.CUDNN, + False, + QKVFusionOption.FULL, + QKVFusionOption.FUSE_KV, + ), + ( + "triton_cudnn_bf16_full_omnidreams_cross", + AttentionBackend.TRITON, + AttentionBackend.OMNIDREAMS, + SDPABackend.CUDNN, + False, + QKVFusionOption.FULL, + QKVFusionOption.NONE, + ), + ) + + native_cases = [case for case in BENCHMARK_CASES if case.native_dit] + assert tuple( + ( + case.implementation, + case.native_dit_backend, + case.native_attention_backend, + case.minimum_compute_capability, + ) + for case in native_cases + ) == ( + ("cuda", "fp8_kvcache_cudnn", "cudnn", None), + ("cuda_sparge", "fp8_kvcache_cudnn", "sparge", None), + ("cuda_sage3", "bf16", "sage3", (12, 0)), + ("cuda_sage3_fp8", "fp8_kvcache_cudnn", "sage3_fp8", (12, 0)), + ) + assert len({case.pytest_id for case in BENCHMARK_CASES}) == len(BENCHMARK_CASES) + + +def test_module_benchmark_cases_cover_attention_policy_matrix() -> None: + """Cover every module SDPA, precision, and QKV fusion combination.""" + triton_cases = [ + case + for case in _MODULE_BENCHMARK_CASES + if case.self_attention_backend is AttentionBackend.TRITON + and case.cross_attention_backend is AttentionBackend.TRITON + ] + expected_cases = { + (sdpa_backend, use_fp8, qkv_fusion_option) + for sdpa_backend in SDPABackend + for use_fp8 in (False, True) + for qkv_fusion_option in QKVFusionOption + } + + assert { + ( + case.sdpa_backend, + case.use_fp8, + case.self_attn_qkv_fusion_option, + ) + for case in triton_cases + } == expected_cases + assert len(_MODULE_BENCHMARK_CASES) == 14 + assert len(_MODULE_SELF_ATTENTION_CASES) == 13 + assert len(_MODULE_CROSS_ATTENTION_CASES) == 9 + + +def test_triton_backend_preserves_checkpoint_keys() -> None: + """Load Omnidreams weights into the Triton block strictly.""" + omnidreams_block = Block(x_dim=32, context_dim=16, num_heads=2) + triton_block = Block( + x_dim=32, + context_dim=16, + num_heads=2, + self_attention_backend=AttentionBackend.TRITON, + cross_attention_backend=AttentionBackend.TRITON, + ) + + triton_block.load_state_dict(omnidreams_block.state_dict(), strict=True) + + +def test_triton_cross_attention_dispatches_tma( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Keep the cross-attention adapter on Triton's backend-owned forward.""" + torch.manual_seed(0) + triton_block = Block( + x_dim=32, + context_dim=16, + num_heads=2, + self_attention_backend=AttentionBackend.TRITON, + cross_attention_backend=AttentionBackend.TRITON, + ) + triton_cross_attention = triton_block.cross_attn + assert isinstance(triton_cross_attention, transformer_modules.TritonCrossAttention) + assert triton_cross_attention.attention_type is AttentionType.CROSS_ATTENTION + assert type(triton_cross_attention).forward is TritonMultiHeadAttention.forward + assert ( + type(triton_cross_attention).compute_kv is TritonMultiHeadAttention.compute_kv + ) + assert ( + type(triton_cross_attention)._attention is TritonMultiHeadAttention._attention + ) + + calls: list[tuple[torch.Size, torch.Size, torch.Size]] = [] + + def record_tma_attention( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + ) -> torch.Tensor: + calls.append((query.shape, key.shape, value.shape)) + return torch.nn.functional.scaled_dot_product_attention( + query.transpose(1, 2), + key.transpose(1, 2), + value.transpose(1, 2), + ).transpose(1, 2) + + monkeypatch.setattr( + triton_attention, + "flash_attention_2_tma", + record_tma_attention, + ) + query = torch.randn(2, 3, 2, 16) + key = torch.randn(2, 5, 2, 16) + value = torch.randn(2, 5, 2, 16) + triton_output = triton_cross_attention._attention(query, key, value) + + assert calls == [ + ( + torch.Size([2, 3, 2, 16]), + torch.Size([2, 5, 2, 16]), + torch.Size([2, 5, 2, 16]), + ) + ] + assert triton_output.shape == query.shape + assert torch.isfinite(triton_output).all() diff --git a/integrations/wan21/README.md b/integrations/wan21/README.md index 895a48e74..a3df2f320 100644 --- a/integrations/wan21/README.md +++ b/integrations/wan21/README.md @@ -113,6 +113,53 @@ video = pipeline.generate(autoregressive_index=0, cache=cache) pipeline.finalize(autoregressive_index=0, cache=cache) # update one-step stats ``` +## Benchmarks + +The WAN21 benchmarks are manual, GPU-only pytest tests for the shipped +`wan21-t2v-1.3b-480p` configuration. Each benchmark layer compares the default +WAN/cuDNN self-attention path, Triton FP8 projections with PyTorch cuDNN SDPA, +and Triton FP8 projections with Triton FlashAttention2 (FA2). Their stable +labels are `wan_torch`, `triton_cudnn`, and `triton_fa2`. Both Triton cases +require an NVIDIA GPU with compute capability 9.0 or newer and do not support +context parallelism. + +First sync this integration and the workspace benchmark dependencies: + +```bash +uv sync --package flashdreams-wan21 --group test +``` + +Run the complete suite from the workspace root: + +```bash +uv run --package flashdreams-wan21 --group test pytest \ + integrations/wan21/benchmarks \ + -p no:manual_marker -m manual --benchmark-only -v +``` + +To run one benchmark layer, replace the benchmark directory with one of these +files: + +- `test_modules.py` benchmarks self-attention and one complete DiT block at the + production 480x832, 21-latent-frame tensor geometry with shared random + weights. +- `test_network.py` benchmarks one complete, random-initialized Wan 2.1 1.3B + DiT evaluation at the same geometry. +- `test_pipeline.py` separately benchmarks checkpoint-backed `generate` and + `finalize` for the production 50-step CFG pipeline. It records one measured + full `generate` per backend and constructs a production-shaped synthetic + final state for `finalize`, so the finalize case runs no denoising rollout. + +The pipeline uses targeted untimed DiT and VAE component prewarm instead of a +full-generation warmup. Its results are single-sample latency and throughput +observations, not median or p90 estimates. The benchmarks exclude model +construction, checkpoint loading, and cache-object allocation from measured +rounds. The one-shot pipeline measurements intentionally include AR=0 +CUDA-graph wrapper input staging, which is recurring per-rollout work after a +fresh cache resets captured graph state. When publishing results, also record +the exact command and commit, prompt and seed, compiler-cache state, GPU/driver +and software stack, and any fallback warnings. + ## Tests ```bash diff --git a/integrations/wan21/benchmarks/cases.py b/integrations/wan21/benchmarks/cases.py new file mode 100644 index 000000000..f572c3675 --- /dev/null +++ b/integrations/wan21/benchmarks/cases.py @@ -0,0 +1,106 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared Wan 2.1 attention benchmark cases.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import pytest +import torch + +from flashdreams.accelerated.multi_head_attention_triton import SDPABackend +from flashdreams.recipes.wan.transformer.impl.modules import AttentionBackend + + +@dataclass(frozen=True) +class AttentionBenchmarkCase: + """Configuration and metadata for one attention benchmark implementation.""" + + implementation: str + """Stable implementation name stored in benchmark metadata.""" + + attention_backend: AttentionBackend + """DiT block implementation configured for this case.""" + + sdpa_backend: SDPABackend + """SDPA implementation configured for Triton self-attention.""" + + self_attention_operator: str + """Self-attention operator reported in benchmark metadata.""" + + cross_attention_operator: str + """Cross-attention operator reported in benchmark metadata.""" + + minimum_compute_capability: tuple[int, int] | None = None + """Minimum CUDA compute capability; ``None`` accepts any CUDA device.""" + + @property + def pytest_id(self) -> str: + """Return the readable pytest parameter identifier.""" + return self.implementation.replace("_", "-") + + +WAN_TORCH_CASE = AttentionBenchmarkCase( + implementation="wan_torch", + attention_backend=AttentionBackend.WAN, + sdpa_backend=SDPABackend.CUDNN, + self_attention_operator="cudnn", + cross_attention_operator="cudnn", +) + +TRITON_CUDNN_CASE = AttentionBenchmarkCase( + implementation="triton_cudnn", + attention_backend=AttentionBackend.TRITON, + sdpa_backend=SDPABackend.CUDNN, + self_attention_operator="torch_cudnn_sdpa", + cross_attention_operator="triton_fa2", + minimum_compute_capability=(9, 0), +) + +TRITON_FA2_CASE = AttentionBenchmarkCase( + implementation="triton_fa2", + attention_backend=AttentionBackend.TRITON, + sdpa_backend=SDPABackend.TRITON, + self_attention_operator="triton_fa2", + cross_attention_operator="triton_fa2", + minimum_compute_capability=(9, 0), +) + +ATTENTION_CASES = (WAN_TORCH_CASE, TRITON_CUDNN_CASE, TRITON_FA2_CASE) +"""Attention cases exercised by each Wan 2.1 benchmark layer.""" + +assert {case.attention_backend for case in ATTENTION_CASES} == set(AttentionBackend) +assert { + case.sdpa_backend + for case in ATTENTION_CASES + if case.attention_backend is AttentionBackend.TRITON +} == set(SDPABackend) + + +def skip_unsupported_device( + case: AttentionBenchmarkCase, + device: torch.device, +) -> None: + """Skip a benchmark case when device is older than its minimum capability.""" + minimum = case.minimum_compute_capability + if minimum is None: + return + if torch.cuda.get_device_capability(device) < minimum: + pytest.skip( + f"{case.pytest_id} attention requires compute capability " + f"{minimum[0]}.{minimum[1]}+" + ) diff --git a/integrations/wan21/benchmarks/test_modules.py b/integrations/wan21/benchmarks/test_modules.py new file mode 100644 index 000000000..79783b017 --- /dev/null +++ b/integrations/wan21/benchmarks/test_modules.py @@ -0,0 +1,536 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Microbenchmarks for Wan 2.1 self-attention and DiT blocks. + +Run all attention cases with:: + + uv run --package flashdreams-wan21 --group test pytest \ + integrations/wan21/benchmarks/test_modules.py \ + -p no:manual_marker -m manual --benchmark-only +""" + +from __future__ import annotations + +import os + +import pytest +import torch +import torch.distributed as dist +from pytest_benchmark.fixture import BenchmarkFixture + +from flashdreams.core.attention.rope import RotaryPositionEmbedding3D +from flashdreams.core.distributed import init as init_distributed +from flashdreams.recipes.wan.transformer.impl.modules import ( + AttentionBackend, + Block, +) +from flashdreams.recipes.wan.transformer.impl.network import ( + WanDiTNetwork1pt3BConfig, +) +from integrations.wan21.benchmarks.cases import ( + ATTENTION_CASES, + AttentionBenchmarkCase, + skip_unsupported_device, +) + +pytestmark = pytest.mark.manual + +_GPU_REASON = "Wan 2.1 DiT module benchmarks require CUDA" + +# The shipped T2V runner generates 480x832 pixels from a single 21-frame +# latent chunk. Wan VAE compression yields 21x60x104 latents, and the DiT's +# 1x2x2 patching produces 21x30x52 attention tokens. +_PIXEL_HEIGHT = 480 +_PIXEL_WIDTH = 832 +_LATENT_HEIGHT = 60 +_LATENT_WIDTH = 104 +_CHUNK_SIZE_T = 21 +_WINDOW_SIZE_T = 21 +_SINK_SIZE_T = 0 +_TEXT_TOKENS = 512 +_WARMUP_ROUNDS = 3 +_BENCHMARK_ROUNDS = 20 +_SEED = 42 + + +def _benchmark_device() -> torch.device: + """Initialize context parallelism and return this rank's GPU.""" + if int(os.environ.get("WORLD_SIZE", "1")) > 1 and not dist.is_initialized(): + init_distributed() + if dist.is_initialized(): + return torch.device("cuda", torch.cuda.current_device()) + torch.cuda.set_device(0) + return torch.device("cuda", 0) + + +def _synchronize_ranks() -> None: + """Align context-parallel ranks before a benchmark sample.""" + if dist.is_initialized(): + dist.barrier() + + +def _skip_unsupported_case( + case: AttentionBenchmarkCase, + device: torch.device, + context_parallel_size: int, +) -> None: + """Skip case and execution combinations unsupported by production code.""" + skip_unsupported_device(case, device) + if case.attention_backend is AttentionBackend.TRITON and context_parallel_size > 1: + pytest.skip("Triton attention does not support context parallelism") + + +def _make_block( + config: WanDiTNetwork1pt3BConfig, + case: AttentionBenchmarkCase, + device: torch.device, + dtype: torch.dtype, +) -> Block: + """Build a backend-selected block with shared random weights.""" + + def make(selected_backend: AttentionBackend) -> Block: + return Block( + dim=config.dim, + ffn_dim=config.ffn_dim, + num_heads=config.num_heads, + cross_attn_norm=config.cross_attn_norm, + eps=config.eps, + i2v=config.cross_attn_enable_img, + apply_rope_before_kvcache=config.apply_rope_before_kvcache, + cp_method=config.cp_method, + attention_backend=selected_backend, + sdpa_backend=config.sdpa_backend, + ) + + # Allocate the block directly at its benchmark precision. Initialization + # and state-dict conversion remain outside the measured region. + previous_dtype = torch.get_default_dtype() + try: + torch.set_default_dtype(dtype) + with torch.device(device): + torch.manual_seed(_SEED) + reference = make(AttentionBackend.WAN) + if case.attention_backend is AttentionBackend.WAN: + return reference + block = make(case.attention_backend) + block.load_state_dict(reference.state_dict(), strict=True) + return block + finally: + torch.set_default_dtype(previous_dtype) + + +def _token_geometry( + config: WanDiTNetwork1pt3BConfig, + context_parallel_size: int, +) -> tuple[int, int, int, int, int, int, int, int]: + """Return global and per-rank token geometry for the shipped T2V shape.""" + patch_t = _CHUNK_SIZE_T // config.patch_size[0] + patch_h = _LATENT_HEIGHT // config.patch_size[1] + patch_w = _LATENT_WIDTH // config.patch_size[2] + tokens_per_frame = patch_h * patch_w + global_chunk_tokens = patch_t * tokens_per_frame + global_window_tokens = _WINDOW_SIZE_T * tokens_per_frame + global_sink_tokens = _SINK_SIZE_T * tokens_per_frame + assert global_chunk_tokens % context_parallel_size == 0 + assert global_window_tokens % context_parallel_size == 0 + assert global_sink_tokens % context_parallel_size == 0 + return ( + patch_t, + patch_h, + patch_w, + global_chunk_tokens, + global_window_tokens, + global_sink_tokens, + global_chunk_tokens // context_parallel_size, + global_window_tokens // context_parallel_size, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason=_GPU_REASON) +@pytest.mark.parametrize( + "case", + ATTENTION_CASES, + ids=lambda case: case.pytest_id, +) +@torch.inference_mode() +def test_self_attention_benchmark( + benchmark: BenchmarkFixture, + case: AttentionBenchmarkCase, +) -> None: + """Benchmark Wan 2.1 self-attention over its full single-chunk window.""" + device = _benchmark_device() + if not torch.cuda.is_bf16_supported(): + pytest.skip("Wan 2.1 self-attention benchmark requires bfloat16 support") + + dtype = torch.bfloat16 + cp_size = dist.get_world_size() if dist.is_initialized() else 1 + _skip_unsupported_case(case, device, cp_size) + backend = case.attention_backend + config = WanDiTNetwork1pt3BConfig( + cp_method="ring", + attention_backend=backend, + sdpa_backend=case.sdpa_backend, + ) + block = _make_block(config, case, device, dtype).eval() + block.update_parameters_after_loading_checkpoint() + attention = block.self_attn + assert block.attention_backend is backend + assert block.sdpa_backend is case.sdpa_backend + del block + + cp_group = dist.group.WORLD if cp_size > 1 else None + attention.set_context_parallel_group(cp_group) + self_attention_cp_enabled = attention.is_context_parallel_enabled() + assert self_attention_cp_enabled == ( + backend is AttentionBackend.WAN and cp_size > 1 + ) + + ( + patch_t, + patch_h, + patch_w, + global_chunk_tokens, + global_window_tokens, + global_sink_tokens, + chunk_tokens, + window_tokens, + ) = _token_geometry(config, cp_size) + sink_tokens = global_sink_tokens // cp_size + head_dim = config.dim // config.num_heads + generator = torch.Generator(device=device).manual_seed(_SEED) + x = torch.randn( + (chunk_tokens, config.dim), + generator=generator, + device=device, + dtype=dtype, + ) + cache = attention.allocate_kv_cache( + batch_size=1, + chunk_size=chunk_tokens, + window_size=window_tokens, + sink_size=sink_tokens, + device=device, + dtype=dtype, + ) + rope = RotaryPositionEmbedding3D( + head_dim=head_dim, + len_t=patch_t, + len_h=patch_h, + len_w=patch_w, + interleaved=True, + device=device, + ) + rope.set_context_parallel_group(cp_group) + rope_freqs = rope.shift_t(0) + + # Populate the one-chunk window before timing. The measured calls overwrite + # the same cache slot, matching repeated denoising evaluations at AR index 0. + cache.before_update(0) + output = attention(x, cache, rope_freqs) + cache.after_update(0) + torch.cuda.synchronize(device) + del output + + benchmark.group = "wan21-dit-self-attention" + benchmark.extra_info.update( + { + "module": "self_attention", + "module_owner": "flashdreams.recipes.wan", + "integration": "wan21", + "model_family": "wan", + "model_variant": "wan21-t2v-1.3b-480p", + "implementation": case.implementation, + "attention_backend": backend.value, + "sdpa_backend": case.sdpa_backend.value, + "self_attention_operator": case.self_attention_operator, + "projection_backend": ( + "separate_qkv" + if backend is AttentionBackend.WAN + else "row_scaled_fp8_fused_qkv_output" + ), + "batch_shape": [], + "flattened_batch_size": 1, + "pixel_resolution": [_PIXEL_HEIGHT, _PIXEL_WIDTH], + "latent_shape": [_CHUNK_SIZE_T, _LATENT_HEIGHT, _LATENT_WIDTH], + "attention_grid": [patch_t, patch_h, patch_w], + "global_chunk_tokens": global_chunk_tokens, + "local_chunk_tokens": chunk_tokens, + "global_window_tokens": global_window_tokens, + "local_window_tokens": window_tokens, + "global_sink_tokens": global_sink_tokens, + "local_sink_tokens": sink_tokens, + "model_channels": config.dim, + "num_heads": config.num_heads, + "head_dim": head_dim, + "parameter_count": sum( + parameter.numel() for parameter in attention.parameters() + ), + "checkpoint": "random_init_shared_weights", + "dtype": str(dtype), + "cache_dtype": str(cache.dtype), + "self_attention_context_parallel_method": ( + config.cp_method if backend is AttentionBackend.WAN else None + ), + "context_parallel_size": cp_size, + "self_attention_context_parallel_enabled": (self_attention_cp_enabled), + "distributed_sample_alignment": "barrier_before_each_round", + "cache_state": "full_single_chunk_window", + "cache_prefill_chunks": 1, + "benchmark_ar_index": 0, + "cache_update_bookkeeping": "excluded_from_timing", + "rope_interleaved": True, + "compiled": False, + "cuda_graph": False, + "global_rank": dist.get_rank() if dist.is_initialized() else 0, + "gpu": torch.cuda.get_device_name(device), + "compute_capability": list(torch.cuda.get_device_capability(device)), + "torch": torch.__version__, + "cuda": torch.version.cuda, + "cudnn": torch.backends.cudnn.version(), + "warmup_rounds": _WARMUP_ROUNDS, + "benchmark_rounds": _BENCHMARK_ROUNDS, + "seed": _SEED, + } + ) + + cache.before_update(0) + torch.cuda.synchronize(device) + torch.cuda.reset_peak_memory_stats(device) + + def synchronized_forward() -> torch.Tensor: + result = attention(x, cache, rope_freqs) + torch.cuda.synchronize(device) + return result + + output = benchmark.pedantic( + synchronized_forward, + setup=_synchronize_ranks, + iterations=1, + rounds=_BENCHMARK_ROUNDS, + warmup_rounds=_WARMUP_ROUNDS, + ) + cache.after_update(0) + benchmark.extra_info["peak_cuda_memory_bytes"] = torch.cuda.max_memory_allocated( + device + ) + + assert output.shape == x.shape + assert torch.isfinite(output).all() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason=_GPU_REASON) +@pytest.mark.parametrize( + "case", + ATTENTION_CASES, + ids=lambda case: case.pytest_id, +) +@torch.inference_mode() +def test_dit_block_benchmark( + benchmark: BenchmarkFixture, + case: AttentionBenchmarkCase, +) -> None: + """Benchmark the Wan 2.1 T2V DiT block over its full attention window.""" + device = _benchmark_device() + if not torch.cuda.is_bf16_supported(): + pytest.skip("Wan 2.1 DiT block benchmark requires bfloat16 support") + + dtype = torch.bfloat16 + cp_size = dist.get_world_size() if dist.is_initialized() else 1 + _skip_unsupported_case(case, device, cp_size) + backend = case.attention_backend + config = WanDiTNetwork1pt3BConfig( + cp_method="ring", + attention_backend=backend, + sdpa_backend=case.sdpa_backend, + ) + block = _make_block(config, case, device, dtype).eval() + block.update_parameters_after_loading_checkpoint() + assert block.attention_backend is backend + assert block.sdpa_backend is case.sdpa_backend + + cp_group = dist.group.WORLD if cp_size > 1 else None + block.set_context_parallel_group(cp_group) + self_attention_cp_enabled = block.self_attn.is_context_parallel_enabled() + cross_attention_cp_enabled = block.cross_attn.is_context_parallel_enabled() + assert self_attention_cp_enabled == ( + backend is AttentionBackend.WAN and cp_size > 1 + ) + assert not cross_attention_cp_enabled + + ( + patch_t, + patch_h, + patch_w, + global_chunk_tokens, + global_window_tokens, + global_sink_tokens, + chunk_tokens, + window_tokens, + ) = _token_geometry(config, cp_size) + sink_tokens = global_sink_tokens // cp_size + head_dim = config.dim // config.num_heads + generator = torch.Generator(device=device).manual_seed(_SEED) + x = torch.randn( + (chunk_tokens, config.dim), + generator=generator, + device=device, + dtype=dtype, + ) + modulation = torch.randn( + (6, config.dim), + generator=generator, + device=device, + dtype=dtype, + ) + context = torch.randn( + (1, _TEXT_TOKENS, config.dim), + generator=generator, + device=device, + dtype=dtype, + ) + cache = block.initialize_cache( + chunk_size=chunk_tokens, + window_size=window_tokens, + sink_size=sink_tokens, + context_text=context, + ) + rope = RotaryPositionEmbedding3D( + head_dim=head_dim, + len_t=patch_t, + len_h=patch_h, + len_w=patch_w, + interleaved=True, + device=device, + ) + rope.set_context_parallel_group(cp_group) + rope_freqs = rope.shift_t(0) + + def forward() -> torch.Tensor: + cache.before_update(0) + result = block( + x=x, + e=modulation, + cache=cache, + rope_freqs=rope_freqs, + ) + cache.after_update(0) + return result + + # Populate the one-chunk window before timing. Cross-attention's projected + # text context remains static throughout all measured denoising evaluations. + output = forward() + torch.cuda.synchronize(device) + del output + + benchmark.group = "wan21-dit-block" + benchmark.extra_info.update( + { + "module": "Block", + "module_owner": "flashdreams.recipes.wan", + "integration": "wan21", + "model_family": "wan", + "model_variant": "wan21-t2v-1.3b-480p", + "benchmark_scope": "whole_wan_t2v_block", + "implementation": case.implementation, + "attention_backend": backend.value, + "sdpa_backend": case.sdpa_backend.value, + "self_attention_operator": case.self_attention_operator, + "cross_attention_operator": case.cross_attention_operator, + "projection_backend": ( + "separate_qkv" + if backend is AttentionBackend.WAN + else "row_scaled_fp8_fused_qkv_output" + ), + "batch_shape": [], + "flattened_batch_size": 1, + "pixel_resolution": [_PIXEL_HEIGHT, _PIXEL_WIDTH], + "latent_shape": [_CHUNK_SIZE_T, _LATENT_HEIGHT, _LATENT_WIDTH], + "attention_grid": [patch_t, patch_h, patch_w], + "global_chunk_tokens": global_chunk_tokens, + "local_chunk_tokens": chunk_tokens, + "global_window_tokens": global_window_tokens, + "local_window_tokens": window_tokens, + "global_sink_tokens": global_sink_tokens, + "local_sink_tokens": sink_tokens, + "text_tokens": _TEXT_TOKENS, + "model_channels": config.dim, + "ffn_channels": config.ffn_dim, + "num_heads": config.num_heads, + "head_dim": head_dim, + "parameter_count": sum( + parameter.numel() for parameter in block.parameters() + ), + "checkpoint": "random_init_shared_weights", + "dtype": str(dtype), + "self_attention_cache_dtype": str(cache.self_attn.dtype), + "cross_attention_cache_dtype": str(cache.cross_attn.text.dtype), + "self_attention_context_parallel_method": ( + config.cp_method if backend is AttentionBackend.WAN else None + ), + "cross_attention_method": ( + config.cp_method if backend is AttentionBackend.WAN else None + ), + "context_parallel_size": cp_size, + "self_attention_context_parallel_enabled": (self_attention_cp_enabled), + "cross_attention_context_parallel_enabled": (cross_attention_cp_enabled), + "distributed_sample_alignment": "barrier_before_each_round", + "cache_state": "full_single_chunk_window_static_text", + "cache_prefill_chunks": 1, + "benchmark_ar_index": 0, + "cache_update_bookkeeping": "excluded_from_timing", + "rope_interleaved": True, + "compiled": False, + "cuda_graph": False, + "global_rank": dist.get_rank() if dist.is_initialized() else 0, + "gpu": torch.cuda.get_device_name(device), + "compute_capability": list(torch.cuda.get_device_capability(device)), + "torch": torch.__version__, + "cuda": torch.version.cuda, + "cudnn": torch.backends.cudnn.version(), + "warmup_rounds": _WARMUP_ROUNDS, + "benchmark_rounds": _BENCHMARK_ROUNDS, + "seed": _SEED, + } + ) + + cache.before_update(0) + torch.cuda.synchronize(device) + torch.cuda.reset_peak_memory_stats(device) + + def synchronized_forward() -> torch.Tensor: + result = block( + x=x, + e=modulation, + cache=cache, + rope_freqs=rope_freqs, + ) + torch.cuda.synchronize(device) + return result + + output = benchmark.pedantic( + synchronized_forward, + setup=_synchronize_ranks, + iterations=1, + rounds=_BENCHMARK_ROUNDS, + warmup_rounds=_WARMUP_ROUNDS, + ) + cache.after_update(0) + benchmark.extra_info["peak_cuda_memory_bytes"] = torch.cuda.max_memory_allocated( + device + ) + + assert output.shape == x.shape + assert torch.isfinite(output).all() diff --git a/integrations/wan21/benchmarks/test_network.py b/integrations/wan21/benchmarks/test_network.py new file mode 100644 index 000000000..879a7bfcd --- /dev/null +++ b/integrations/wan21/benchmarks/test_network.py @@ -0,0 +1,350 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Benchmark the complete Wan 2.1 T2V 1.3B DiT network by backend.""" + +from __future__ import annotations + +import os + +import pytest +import torch +import torch.distributed as dist +from pytest_benchmark.fixture import BenchmarkFixture + +from flashdreams.core.attention import ContextParallelAttention +from flashdreams.core.attention.rope import RotaryPositionEmbedding3D +from flashdreams.core.distributed import init as init_distributed +from flashdreams.infra.acceleration import ( + CUDAGraphDispatch, + cuda_graph_capture_ar_index, +) +from flashdreams.infra.compile import compile_module +from flashdreams.recipes.wan.transformer.impl.modules import AttentionBackend +from flashdreams.recipes.wan.transformer.impl.network import ( + WanDiTNetwork, + WanDiTNetwork1pt3BConfig, +) +from integrations.wan21.benchmarks.cases import ( + ATTENTION_CASES, + AttentionBenchmarkCase, + skip_unsupported_device, +) + +pytestmark = pytest.mark.manual + +_GPU_REASON = "Wan 2.1 DiT network benchmark requires CUDA" + +# Shipped wan21-t2v-1.3b-480p geometry. The 480x832 output becomes a +# 60x104 latent; 1x2x2 patching yields 30x52 tokens for each of 21 frames. +_PIXEL_HEIGHT = 480 +_PIXEL_WIDTH = 832 +_LATENT_HEIGHT = 60 +_LATENT_WIDTH = 104 +_CHUNK_SIZE_T = 21 +_WINDOW_SIZE_T = 21 +_SINK_SIZE_T = 0 +_ATTENTION_HEIGHT = 30 +_ATTENTION_WIDTH = 52 +_TEXT_TOKENS = 512 +_GLOBAL_CHUNK_TOKENS = _CHUNK_SIZE_T * _ATTENTION_HEIGHT * _ATTENTION_WIDTH +_GLOBAL_WINDOW_TOKENS = _WINDOW_SIZE_T * _ATTENTION_HEIGHT * _ATTENTION_WIDTH +_DIFFUSION_TIMESTEP = 1000.0 +_PIPELINE_GUIDANCE_SCALE = 6.0 +_AUTOTUNE_DRAIN_ROUNDS = 3 +_CUDA_GRAPH_WARMUP_ITERS = 2 +_WARMUP_ROUNDS = 3 +_BENCHMARK_ROUNDS = 20 +_SEED = 42 + + +def _benchmark_device() -> torch.device: + """Initialize context parallelism and return this rank's GPU.""" + if int(os.environ.get("WORLD_SIZE", "1")) > 1 and not dist.is_initialized(): + init_distributed() + if dist.is_initialized(): + return torch.device("cuda", torch.cuda.current_device()) + torch.cuda.set_device(0) + return torch.device("cuda", 0) + + +def _synchronize_ranks() -> None: + """Align context-parallel ranks before a benchmark sample.""" + if dist.is_initialized(): + dist.barrier() + + +def _skip_unsupported_case( + case: AttentionBenchmarkCase, + device: torch.device, +) -> None: + """Skip a case where its hardware or execution mode is unsupported.""" + skip_unsupported_device(case, device) + if case.attention_backend is not AttentionBackend.TRITON: + return + if dist.is_initialized() and dist.get_world_size() > 1: + pytest.skip("Triton attention does not support context parallelism") + + +@pytest.mark.parametrize( + "case", + ATTENTION_CASES, + ids=lambda case: case.pytest_id, +) +@pytest.mark.skipif(not torch.cuda.is_available(), reason=_GPU_REASON) +@torch.inference_mode() +def test_dit_network_benchmark( + benchmark: BenchmarkFixture, + case: AttentionBenchmarkCase, +) -> None: + """Benchmark one compiled production-size Wan 2.1 DiT evaluation.""" + device = _benchmark_device() + if not torch.cuda.is_bf16_supported(): + pytest.skip("Wan 2.1 DiT network benchmark requires bfloat16 support") + + dtype = torch.bfloat16 + torch.manual_seed(_SEED) + _skip_unsupported_case(case, device) + backend = case.attention_backend + config = WanDiTNetwork1pt3BConfig( + patch_embedding_type="conv3d", + cp_method="ring", + attention_backend=backend, + sdpa_backend=case.sdpa_backend, + ) + + # Avoid materializing the 1.3B random initialization as fp32 CPU weights. + previous_dtype = torch.get_default_dtype() + try: + torch.set_default_dtype(dtype) + with torch.device(device): + network = WanDiTNetwork(config) + finally: + torch.set_default_dtype(previous_dtype) + network.eval() + network.update_parameters_after_loading_checkpoint() + parameter_count = sum(parameter.numel() for parameter in network.parameters()) + + cp_size = dist.get_world_size() if dist.is_initialized() else 1 + cp_group = dist.group.WORLD if cp_size > 1 else None + network.set_context_parallel_group(cp_group) + assert all( + block.attention_backend is backend and block.sdpa_backend is case.sdpa_backend + for block in network.blocks + ) + attention_modules = [ + module + for module in network.modules() + if isinstance(module, ContextParallelAttention) + ] + assert {attention.backend for attention in attention_modules} == ( + {"cudnn"} if backend is AttentionBackend.WAN else set() + ) + cp_enabled_attention_modules = [ + attention + for attention in attention_modules + if attention.is_context_parallel_enabled() + ] + local_attention_methods = { + attention.method + for attention in attention_modules + if not attention.is_context_parallel_enabled() + } + assert all( + attention.context_parallel_size() == cp_size + for attention in cp_enabled_attention_modules + ) + assert all( + attention.method == config.cp_method + for attention in cp_enabled_attention_modules + ) + assert bool(cp_enabled_attention_modules) == (cp_size > 1) + + assert _GLOBAL_CHUNK_TOKENS % cp_size == 0 + assert _GLOBAL_WINDOW_TOKENS % cp_size == 0 + chunk_tokens = _GLOBAL_CHUNK_TOKENS // cp_size + window_tokens = _GLOBAL_WINDOW_TOKENS // cp_size + patch_volume = config.patch_size[0] * config.patch_size[1] * config.patch_size[2] + generator = torch.Generator(device=device).manual_seed(_SEED) + x = torch.randn( + (chunk_tokens, config.in_dim * patch_volume), + generator=generator, + device=device, + dtype=dtype, + ) + timestep = torch.tensor(_DIFFUSION_TIMESTEP, device=device, dtype=dtype) + text_embeddings = torch.randn( + (_TEXT_TOKENS, config.text_dim), + generator=generator, + device=device, + dtype=dtype, + ) + cache = network.initialize_cache( + chunk_size=chunk_tokens, + window_size=window_tokens, + sink_size=0, + text_embeddings=text_embeddings, + ) + rope = RotaryPositionEmbedding3D( + head_dim=config.dim // config.num_heads, + len_t=_CHUNK_SIZE_T, + len_h=_ATTENTION_HEIGHT, + len_w=_ATTENTION_WIDTH, + interleaved=True, + device=device, + ) + rope.set_context_parallel_group(cp_group) + rope_freqs = rope.shift_t(0) + + network = compile_module(network) + capture_ar_index = cuda_graph_capture_ar_index( + sink_size_t=_SINK_SIZE_T, + window_size_t=_WINDOW_SIZE_T, + len_t=_CHUNK_SIZE_T, + ) + graph_dispatch = CUDAGraphDispatch( + network, + enabled=True, + capture_ar_idx=capture_ar_index, + warmup_iters=_CUDA_GRAPH_WARMUP_ITERS, + ) + + def forward() -> torch.Tensor: + # The shipped integration generates only AR index 0. Its window is one + # full chunk, so capture starts at index 1 and production uses drain. + return graph_dispatch.select(0, uncond=False)( + x=x, + timesteps=timestep, + cache=cache, + rope_freqs=rope_freqs, + current_chunk_idx=0, + eager_mode=False, + ) + + assert capture_ar_index == 1 + cache.before_update(0) + for _ in range(_AUTOTUNE_DRAIN_ROUNDS): + output = forward() + torch.cuda.synchronize(device) + del output + + benchmark.group = "wan21-t2v-1.3b-dit-network" + benchmark.extra_info.update( + { + "network": "WanDiTNetwork1pt3B", + "integration": "wan21", + "model_family": "wan", + "model_variant": "wan21-t2v-1.3b-480p", + "implementation": case.implementation, + "execution_backend": "pytorch", + "attention_backend": backend.value, + "sdpa_backend": case.sdpa_backend.value, + "self_attention_operator": case.self_attention_operator, + "cross_attention_operator": case.cross_attention_operator, + "projection_backend": ( + "separate_qkv" + if backend is AttentionBackend.WAN + else "row_scaled_fp8_fused_qkv_output" + ), + "batch_shape": [], + "flattened_batch_size": 1, + "pixel_resolution": [_PIXEL_HEIGHT, _PIXEL_WIDTH], + "latent_shape": [_CHUNK_SIZE_T, _LATENT_HEIGHT, _LATENT_WIDTH], + "attention_grid": [ + _CHUNK_SIZE_T, + _ATTENTION_HEIGHT, + _ATTENTION_WIDTH, + ], + "global_chunk_tokens": _GLOBAL_CHUNK_TOKENS, + "local_chunk_tokens": chunk_tokens, + "global_window_tokens": _GLOBAL_WINDOW_TOKENS, + "local_window_tokens": window_tokens, + "global_sink_tokens": 0, + "local_sink_tokens": 0, + "text_tokens": _TEXT_TOKENS, + "input_patch_channels": config.in_dim * patch_volume, + "output_patch_channels": config.out_dim * patch_volume, + "model_channels": config.dim, + "ffn_channels": config.ffn_dim, + "num_blocks": config.num_layers, + "num_heads": config.num_heads, + "head_dim": config.dim // config.num_heads, + "parameter_count": parameter_count, + "checkpoint": "random_init_seed_matched", + "dtype": str(dtype), + "self_attention_cache_dtype": str(cache[0].self_attn.dtype), + "cross_attention_cache_dtype": str(cache[0].cross_attn.text.dtype), + "self_attention_context_parallel_method": ( + config.cp_method if backend is AttentionBackend.WAN else None + ), + "local_attention_methods": sorted(local_attention_methods), + "context_parallel_size": cp_size, + "context_parallel_attention_modules": len(cp_enabled_attention_modules), + "local_attention_modules": ( + len(attention_modules) - len(cp_enabled_attention_modules) + ), + "distributed_sample_alignment": "barrier_before_each_round", + "compiled": True, + "compile_mode": "max-autotune-no-cudagraphs", + "cuda_graph_configured": True, + "cuda_graph_selected": False, + "cuda_graph_dispatch": "drain", + "cuda_graph_capture_ar_index": capture_ar_index, + "cuda_graph_warmup_iters": _CUDA_GRAPH_WARMUP_ITERS, + "cache_state": "single_full_chunk_repeated_scheduler_slot", + "benchmark_ar_index": 0, + "cache_update_bookkeeping": "excluded_from_timing", + "diffusion_timestep": _DIFFUSION_TIMESTEP, + "pipeline_guidance_scale": _PIPELINE_GUIDANCE_SCALE, + "classifier_free_guidance_branch": "conditional", + "rope_interleaved": True, + "global_rank": dist.get_rank() if dist.is_initialized() else 0, + "gpu": torch.cuda.get_device_name(device), + "compute_capability": list(torch.cuda.get_device_capability(device)), + "torch": torch.__version__, + "cuda": torch.version.cuda, + "cudnn": torch.backends.cudnn.version(), + "autotune_drain_rounds": _AUTOTUNE_DRAIN_ROUNDS, + "warmup_rounds": _WARMUP_ROUNDS, + "benchmark_rounds": _BENCHMARK_ROUNDS, + "compiler_cache_state": ( + "host-dependent; compile and autotune excluded from measured rounds" + ), + "seed": _SEED, + } + ) + + torch.cuda.reset_peak_memory_stats(device) + + def synchronized_forward() -> torch.Tensor: + result = forward() + torch.cuda.synchronize(device) + return result + + output = benchmark.pedantic( + synchronized_forward, + setup=_synchronize_ranks, + iterations=1, + rounds=_BENCHMARK_ROUNDS, + warmup_rounds=_WARMUP_ROUNDS, + ) + cache.after_update(0) + benchmark.extra_info["peak_cuda_memory_bytes"] = torch.cuda.max_memory_allocated( + device + ) + + expected_output_shape = (chunk_tokens, config.out_dim * patch_volume) + assert output.shape == expected_output_shape + assert torch.isfinite(output).all() diff --git a/integrations/wan21/benchmarks/test_pipeline.py b/integrations/wan21/benchmarks/test_pipeline.py new file mode 100644 index 000000000..7f41b57b2 --- /dev/null +++ b/integrations/wan21/benchmarks/test_pipeline.py @@ -0,0 +1,594 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""One-shot full-pipeline benchmarks for the shipped Wan 2.1 T2V runner. + +Run the manual GPU benchmarks with:: + + uv run --package flashdreams-wan21 --group test pytest \ + integrations/wan21/benchmarks/test_pipeline.py \ + -p no:manual_marker -m manual --benchmark-only -v +""" + +from __future__ import annotations + +import os +from typing import Literal + +import pytest +import torch +import torch.distributed as dist +from pytest_benchmark.fixture import BenchmarkFixture +from wan21.config import PIPELINE_WAN21_T2V_1PT3B_480P +from wan21.runner import DEFAULT_PROMPT + +from flashdreams.core.attention import ContextParallelAttention +from flashdreams.core.distributed import init as init_distributed +from flashdreams.infra.config import derive_config +from flashdreams.infra.diffusion.model import DiffusionModel +from flashdreams.infra.diffusion.scheduler import ( + FlowMatchUniPCScheduler, + FlowMatchUniPCSchedulerConfig, +) +from flashdreams.infra.encoder.text.umt5 import UMT5TextEncoderConfig +from flashdreams.infra.pipeline import StreamInferencePipeline +from flashdreams.recipes.wan import ( + NEGATIVE_PROMPT, + Wan21Transformer, + Wan21TransformerConfig, + WanDiTNetwork, + WanInferencePipeline, + WanInferencePipelineCache, + WanVAEDecoderConfig, +) +from flashdreams.recipes.wan.transformer.impl.modules import AttentionBackend +from flashdreams.recipes.wan.transformer.wan21 import Wan21TransformerCache +from integrations.wan21.benchmarks.cases import ( + ATTENTION_CASES, + AttentionBenchmarkCase, + skip_unsupported_device, +) + +pytestmark = pytest.mark.manual + +_GPU_REASON = "Wan 2.1 full-pipeline benchmark requires CUDA" + +_PIXEL_HEIGHT = 480 +_PIXEL_WIDTH = 832 +_TEXT_TOKENS = 512 +_COMPONENT_PREWARM_ROUNDS = 3 +_WARMUP_ROUNDS = 0 +_BENCHMARK_ROUNDS = 1 +_SEED = 42 + + +def _benchmark_device() -> torch.device: + """Initialize context parallelism and return this rank's GPU.""" + if int(os.environ.get("WORLD_SIZE", "1")) > 1 and not dist.is_initialized(): + init_distributed() + if dist.is_initialized(): + return torch.device("cuda", torch.cuda.current_device()) + torch.cuda.set_device(0) + return torch.device("cuda", 0) + + +def _synchronize_ranks() -> None: + """Align context-parallel ranks outside a benchmark sample.""" + if dist.is_initialized(): + dist.barrier() + + +def _skip_unsupported_case( + case: AttentionBenchmarkCase, + device: torch.device, +) -> None: + """Skip a case where its hardware or context-parallel contract is unmet.""" + skip_unsupported_device(case, device) + if case.attention_backend is not AttentionBackend.TRITON: + return + world_size = ( + dist.get_world_size() + if dist.is_initialized() + else int(os.environ.get("WORLD_SIZE", "1")) + ) + if world_size > 1: + pytest.skip("Triton attention does not support context parallelism") + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason=_GPU_REASON) +@pytest.mark.parametrize( + "case", + ATTENTION_CASES, + ids=lambda case: case.pytest_id, +) +def test_full_pipeline_generate_benchmark( + benchmark: BenchmarkFixture, + case: AttentionBenchmarkCase, +) -> None: + """Benchmark the shipped one-shot Wan 2.1 denoise and decode path.""" + _run_full_pipeline_benchmark(benchmark, case=case, stage="generate") + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason=_GPU_REASON) +@pytest.mark.parametrize( + "case", + ATTENTION_CASES, + ids=lambda case: case.pytest_id, +) +def test_full_pipeline_finalize_benchmark( + benchmark: BenchmarkFixture, + case: AttentionBenchmarkCase, +) -> None: + """Benchmark the matching AR=0 DiT cache-finalization update.""" + _run_full_pipeline_benchmark(benchmark, case=case, stage="finalize") + + +@torch.inference_mode() +def _run_full_pipeline_benchmark( + benchmark: BenchmarkFixture, + *, + case: AttentionBenchmarkCase, + stage: Literal["generate", "finalize"], +) -> None: + """Run one backend and one production pipeline lifecycle stage.""" + device = _benchmark_device() + if not torch.cuda.is_bf16_supported(): + pytest.skip("Wan 2.1 full-pipeline benchmark requires bfloat16 support") + _skip_unsupported_case(case, device) + + torch.manual_seed(_SEED) + torch.backends.cudnn.benchmark = True + + # UMT5 is a one-shot rollout initializer. Use correctly shaped synthetic + # embeddings so setup measures the checkpoint-backed recurring pipeline: + # 50-step CFG diffusion followed by the production Wan VAE decoder. + pipeline_config = derive_config( + PIPELINE_WAN21_T2V_1PT3B_480P, + name=f"wan21-t2v-1.3b-full-pipeline-{case.implementation}-benchmark", + text_encoder=None, + enable_sync_and_profile=False, + diffusion_model={ + "seed": _SEED, + "transformer": { + "init_device": str(device), + "network": { + "attention_backend": case.attention_backend, + "sdpa_backend": case.sdpa_backend, + }, + }, + }, + ) + pipeline = pipeline_config.setup().to(device=device) + assert isinstance(pipeline, WanInferencePipeline) + pipeline.eval() + assert pipeline.encoder is None + assert pipeline.decoder is not None + + recurring_parameter_count = sum( + parameter.numel() for parameter in pipeline.parameters() + ) + source_text_encoder_config = PIPELINE_WAN21_T2V_1PT3B_480P.text_encoder + assert isinstance(source_text_encoder_config, UMT5TextEncoderConfig) + + diffusion_config = pipeline_config.diffusion_model + transformer_config = diffusion_config.transformer + scheduler_config = diffusion_config.scheduler + decoder_config = pipeline_config.decoder + assert isinstance(transformer_config, Wan21TransformerConfig) + assert isinstance(scheduler_config, FlowMatchUniPCSchedulerConfig) + assert isinstance(decoder_config, WanVAEDecoderConfig) + assert transformer_config.network.attention_backend is case.attention_backend + assert transformer_config.network.sdpa_backend is case.sdpa_backend + assert transformer_config.batch_shape == () + assert transformer_config.len_t == 21 + assert transformer_config.window_size_t == 21 + assert transformer_config.sink_size_t == 0 + assert transformer_config.guidance_scale == 6.0 + assert scheduler_config.num_inference_steps == 50 + + transformer = pipeline.diffusion_model.transformer + scheduler = pipeline.diffusion_model.scheduler + assert isinstance(transformer, Wan21Transformer) + assert isinstance(scheduler, FlowMatchUniPCScheduler) + network = getattr(transformer.network, "_orig_mod", transformer.network) + assert isinstance(network, WanDiTNetwork) + assert network.blocks + assert all( + block.attention_backend is case.attention_backend + and block.sdpa_backend is case.sdpa_backend + for block in network.blocks + ) + + context_parallel_attention_modules = [ + module + for module in pipeline.modules() + if isinstance(module, ContextParallelAttention) + ] + assert {attention.backend for attention in context_parallel_attention_modules} == ( + {"cudnn"} if case.attention_backend is AttentionBackend.WAN else set() + ) + cp_size = transformer._cp_size + cp_enabled_attention_modules = [ + attention + for attention in context_parallel_attention_modules + if attention.is_context_parallel_enabled() + ] + local_attention_methods = { + attention.method + for attention in context_parallel_attention_modules + if not attention.is_context_parallel_enabled() + } + assert all( + attention.context_parallel_size() == cp_size + for attention in cp_enabled_attention_modules + ) + assert all( + attention.method == transformer_config.network.cp_method + for attention in cp_enabled_attention_modules + ) + assert bool(cp_enabled_attention_modules) == (cp_size > 1) + + dtype = transformer_config.dtype + spatial_compression = int(pipeline.decoder.spatial_compression_ratio) + latent_height = _PIXEL_HEIGHT // spatial_compression + latent_width = _PIXEL_WIDTH // spatial_compression + latent_channels = int(transformer_config.network.out_dim) + text_dim = int(transformer_config.network.text_dim) + output_frames = pipeline.decoder.get_output_temporal_size( + 0, transformer_config.len_t + ) + patch_t, patch_h, patch_w = transformer_config.network.patch_size + global_tokens = ( + (transformer_config.len_t // patch_t) + * (latent_height // patch_h) + * (latent_width // patch_w) + ) + + text_embeddings = torch.zeros( + (1, _TEXT_TOKENS, text_dim), + device=device, + dtype=dtype, + ) + negative_text_embeddings = torch.zeros_like(text_embeddings) + + def initialize_cache() -> WanInferencePipelineCache: + """Build one production-shaped T2V cache outside sample timing.""" + parent_cache = StreamInferencePipeline.initialize_cache( + pipeline, + transformer_context={ + "height": latent_height, + "width": latent_width, + "text_embeddings": text_embeddings, + "negative_text_embeddings": negative_text_embeddings, + "image_embeddings": None, + }, + ) + cache = WanInferencePipelineCache( + transformer_cache=parent_cache.transformer_cache, + encoder_cache=parent_cache.encoder_cache, + decoder_cache=parent_cache.decoder_cache, + image=None, + ) + assert isinstance(cache.transformer_cache, Wan21TransformerCache) + assert cache.transformer_cache.network_cache_uncond is not None + return cache + + # Prime lazy compilation and cuDNN selection without paying for another + # complete 50-step rollout. Fresh measured caches reset the graph wrappers, + # so their per-rollout AR0 input staging remains inside the timed call. + probe_cache = initialize_cache() + assert isinstance(probe_cache.transformer_cache, Wan21TransformerCache) + first_block_cache = probe_cache.transformer_cache.network_cache.block_caches[0] + self_attention_cache_dtype = str(first_block_cache.self_attn.dtype) + cross_attention_cache_dtype = str(first_block_cache.cross_attn.text.dtype) + + probe_transformer_cache = probe_cache.transformer_cache + probe_transformer_cache.start(0) + prewarm_latent = torch.zeros( + transformer.latent_shape, + device=device, + dtype=dtype, + ) + prewarm_timestep = scheduler.timesteps[0].to(device=device, dtype=dtype) + prewarm_flow = transformer.predict_flow( + noisy_latent=prewarm_latent, + timestep=prewarm_timestep, + cache=probe_transformer_cache, + ) + for _ in range(1, _COMPONENT_PREWARM_ROUNDS): + prewarm_flow = transformer.predict_flow( + noisy_latent=prewarm_latent, + timestep=prewarm_timestep, + cache=probe_transformer_cache, + ) + probe_transformer_cache.finalize(0) + + decoder_prewarm_calls = 0 + prewarm_decode_input: torch.Tensor | None = None + prewarm_decode_output: torch.Tensor | None = None + if stage == "generate": + assert probe_cache.decoder_cache is not None + prewarm_decode_input = torch.zeros( + ( + transformer_config.len_t, + latent_channels, + latent_height, + latent_width, + ), + device=device, + dtype=dtype, + ) + prewarm_decode_output = pipeline.decoder( + input=prewarm_decode_input, + autoregressive_index=0, + cache=probe_cache.decoder_cache, + ) + decoder_prewarm_calls = 1 + + torch.cuda.synchronize(device) + del prewarm_flow, prewarm_latent, prewarm_timestep + prewarm_decode_input = None + del prewarm_decode_output + probe_cache = None + + resolved_timesteps = scheduler.timesteps.detach().cpu().tolist() + resolved_sigmas = scheduler.sigmas.detach().cpu().tolist() + benchmark.group = f"wan21-full-pipeline-{stage}" + benchmark.extra_info.update( + { + "pipeline": pipeline_config.name, + "source_pipeline": PIPELINE_WAN21_T2V_1PT3B_480P.name, + "integration": "wan21", + "model_family": "wan", + "model_variant": "wan21-t2v-1.3b-480p", + "batch_shape": list(transformer_config.batch_shape), + "pixel_resolution": [_PIXEL_HEIGHT, _PIXEL_WIDTH], + "latent_shape": [ + transformer_config.len_t, + latent_channels, + latent_height, + latent_width, + ], + "attention_grid": [ + transformer_config.len_t // patch_t, + latent_height // patch_h, + latent_width // patch_w, + ], + "global_tokens": global_tokens, + "local_tokens": transformer.latent_shape[-2], + "output_frames": output_frames, + "output_fps": 16, + "autoregressive_index": 0, + "rollout_chunks": 1, + "text_tokens": _TEXT_TOKENS, + "text_embedding_dim": text_dim, + "prompt": DEFAULT_PROMPT, + "negative_prompt": NEGATIVE_PROMPT, + "prompt_embedding_source": "synthetic_precomputed_zeros", + "raw_prompt_encoding_timed": False, + "text_encoder_model": source_text_encoder_config.model_id_or_local_path, + "num_inference_steps": scheduler_config.num_inference_steps, + "scheduler": type(scheduler).__name__, + "scheduler_shift": scheduler_config.shift, + "scheduler_solver_order": scheduler_config.solver_order, + "resolved_timesteps": resolved_timesteps, + "resolved_sigmas_fp32": resolved_sigmas, + "guidance_scale": transformer_config.guidance_scale, + "cfg_network_branches": 2, + "context_noise": diffusion_config.context_noise, + "window_size_t": transformer_config.window_size_t, + "sink_size_t": transformer_config.sink_size_t, + "timed_stage": stage, + "timed_stages": ( + ["unipc_diffuse", "wan_vae_decode"] + if stage == "generate" + else ["dit_cache_finalize"] + ), + "untimed_lifecycle_stage": ( + "finalize" if stage == "generate" else "synthetic_final_state_setup" + ), + "full_generate_calls": 1 if stage == "generate" else 0, + "finalize_state_source": ( + "generated" if stage == "generate" else "synthetic_zero_clean_latent" + ), + "cache_initialization": ( + "cache object allocation excluded; AR0 CUDA graph wrapper " + "staging included" + ), + "checkpoint_loading": "excluded_from_timing", + "dit_checkpoint": transformer_config.checkpoint_path, + "decoder_checkpoint": decoder_config.checkpoint_path, + "dtype": str(dtype), + "implementation": case.implementation, + "dit_execution": "pytorch", + "configured_attention_backend": case.attention_backend.value, + "configured_sdpa_backend": case.sdpa_backend.value, + "dit_self_attention_backend": case.self_attention_operator, + "dit_cross_attention_backend": case.cross_attention_operator, + "projection_backend": ( + "separate_qkv" + if case.attention_backend is AttentionBackend.WAN + else "row_scaled_fp8_fused_qkv_output" + ), + "dit_self_attention_kv_cache_dtype": self_attention_cache_dtype, + "dit_cross_attention_kv_cache_dtype": cross_attention_cache_dtype, + "self_attention_context_parallel_method": ( + transformer_config.network.cp_method + if case.attention_backend is AttentionBackend.WAN + else None + ), + "local_attention_methods": sorted(local_attention_methods), + "context_parallel_size": cp_size, + "context_parallel_attention_modules": len(cp_enabled_attention_modules), + "local_attention_modules": ( + len(context_parallel_attention_modules) + - len(cp_enabled_attention_modules) + ), + "distributed_sample_alignment": "barrier_during_untimed_setup", + "dit_compiled": transformer_config.compile_network, + "dit_cuda_graph": transformer_config.use_cuda_graph, + "dit_cuda_graph_capture_ar_index": transformer._cuda_graph_capture_ar_idx, + "dit_ar0_execution": "eager_drain_before_steady_state_capture", + "decoder_compiled": decoder_config.use_compile, + "decoder_cuda_graph": decoder_config.use_cuda_graph, + "decoder_ar0_execution": "eager_drain_with_fresh_cache", + "recurring_pipeline_parameter_count": recurring_parameter_count, + "global_rank": dist.get_rank() if dist.is_initialized() else 0, + "gpu": torch.cuda.get_device_name(device), + "compute_capability": list(torch.cuda.get_device_capability(device)), + "torch": torch.__version__, + "cuda": torch.version.cuda, + "cudnn": torch.backends.cudnn.version(), + "cudnn_benchmark": torch.backends.cudnn.benchmark, + "targeted_dit_prewarm_calls": _COMPONENT_PREWARM_ROUNDS, + "targeted_decoder_prewarm_calls": decoder_prewarm_calls, + "warmup_rounds": _WARMUP_ROUNDS, + "latency_summary": "single_sample_no_percentiles", + "benchmark_rounds": _BENCHMARK_ROUNDS, + "startup_timing": ( + "model and checkpoint setup excluded; per-rollout AR0 CUDA " + "graph wrapper staging included" + ), + "compiler_cache_state": ( + "checkpoint loading and targeted component prewarm excluded; " + "no full-generate warmup" + ), + "num_gpus_visible": torch.cuda.device_count(), + "seed": _SEED, + "rng_reset_per_sample": True, + } + ) + + cache: WanInferencePipelineCache | None = None + latest_output: torch.Tensor | None = None + stage_peak_cuda_memory_bytes = 0 + + def prepare_sample() -> None: + nonlocal cache, latest_output + _synchronize_ranks() + # Drop the prior cache before allocating the next multi-GiB KV cache, + # allowing the CUDA allocator to reuse its storage without a peak at 2x. + transformer._cuda_graph_dispatch.reset() + latest_output = None + cache = None + cache = initialize_cache() + rng = pipeline.diffusion_model.rng + assert rng is not None + rng.manual_seed(_SEED) + torch.cuda.synchronize(device) + _synchronize_ranks() + + def record_stage_peak_memory() -> None: + nonlocal stage_peak_cuda_memory_bytes + stage_peak_cuda_memory_bytes = max( + stage_peak_cuda_memory_bytes, + int(torch.cuda.max_memory_allocated(device)), + ) + + if stage == "generate": + + def setup_generate() -> None: + prepare_sample() + torch.cuda.reset_peak_memory_stats(device) + + def synchronized_generate() -> torch.Tensor: + nonlocal latest_output + assert cache is not None + latest_output = pipeline.generate(autoregressive_index=0, cache=cache) + torch.cuda.synchronize(device) + return latest_output + + def teardown_generate() -> None: + assert cache is not None + record_stage_peak_memory() + pipeline.finalize(autoregressive_index=0, cache=cache) + torch.cuda.synchronize(device) + + output = benchmark.pedantic( + synchronized_generate, + setup=setup_generate, + teardown=teardown_generate, + iterations=1, + rounds=_BENCHMARK_ROUNDS, + warmup_rounds=_WARMUP_ROUNDS, + ) + else: + + def setup_finalize() -> None: + prepare_sample() + assert cache is not None + assert isinstance(cache.transformer_cache, Wan21TransformerCache) + # AR0 replaces the full one-chunk KV window and context_noise is + # zero, so values do not change which finalize kernels execute. + cache.autoregressive_index = 0 + cache.transformer_cache.start(0) + cache.final_state = DiffusionModel.FinalState( + clean_latent=torch.zeros( + transformer.latent_shape, + device=device, + dtype=dtype, + ), + autoregressive_index=0, + cache=cache.transformer_cache, + ) + torch.cuda.synchronize(device) + _synchronize_ranks() + torch.cuda.reset_peak_memory_stats(device) + + def synchronized_finalize() -> None: + assert cache is not None + pipeline.finalize(autoregressive_index=0, cache=cache) + torch.cuda.synchronize(device) + + def teardown_finalize() -> None: + record_stage_peak_memory() + + benchmark.pedantic( + synchronized_finalize, + setup=setup_finalize, + teardown=teardown_finalize, + iterations=1, + rounds=_BENCHMARK_ROUNDS, + warmup_rounds=_WARMUP_ROUNDS, + ) + output = None + + benchmark.extra_info["peak_cuda_memory_bytes"] = stage_peak_cuda_memory_bytes + assert benchmark.stats is not None + single_sample_stage_s = benchmark.stats.stats.median + benchmark.extra_info.update( + { + f"single_sample_{stage}_ms": single_sample_stage_s * 1_000, + f"single_sample_{stage}_rollouts_per_second": (1.0 / single_sample_stage_s), + } + ) + if stage == "generate": + benchmark.extra_info["single_sample_generate_output_fps"] = ( + output_frames / single_sample_stage_s + ) + + if stage == "generate": + assert output is not None + assert output.shape == ( + output_frames, + 3, + _PIXEL_HEIGHT, + _PIXEL_WIDTH, + ) + assert torch.isfinite(output).all() + else: + assert cache is not None + assert cache.final_state is not None + assert torch.isfinite(cache.final_state.clean_latent).all() diff --git a/pyproject.toml b/pyproject.toml index 02f62aa87..3dae2f753 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -42,6 +42,8 @@ override-dependencies = [ no-build-isolation-package = ["transformer-engine-torch"] [tool.pyright] +exclude = ["flashdreams/flashdreams/accelerated/triton/**"] +ignore = ["flashdreams/flashdreams/accelerated/triton/**"] extraPaths = [ "flashdreams", "apps", @@ -112,7 +114,7 @@ invalid-method-override = "ignore" replace-imports-with-any = ["flash_attn.**", "transformer_engine.**", "triton.**"] [tool.pytest.ini_options] -addopts = "--import-mode=importlib -p flashdreams._pytest_plugins.marker_enforcement" +addopts = "--import-mode=importlib -p flashdreams._pytest_plugins.marker_enforcement --benchmark-time-unit=ms" norecursedirs = [ "parity_check", "parity_check_v2", @@ -131,6 +133,7 @@ markers = [ # flashdreams/_pytest_plugins. test = [ "pytest>=8.0", + "pytest-benchmark>=5.1", "pytest-asyncio>=0.23", "pytest-manual-marker>=2.0", "tomli>=2.0", @@ -138,6 +141,7 @@ test = [ ] lint = [ "pre-commit>=4.3.0", + "ruff==0.12.7", "sphinx>=7.0", "ty>=0.0.39", {include-group = "test"}, @@ -171,3 +175,6 @@ docs-ci = [ "tqdm>=4.60", "transformers>=5.0,<6", ] +dev = [ + "python-dotenv>=1.2.2", +] diff --git a/scripts/benchmark/flashdreams/accelerated/plot.py b/scripts/benchmark/flashdreams/accelerated/plot.py new file mode 100755 index 000000000..fcd013173 --- /dev/null +++ b/scripts/benchmark/flashdreams/accelerated/plot.py @@ -0,0 +1,312 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Performance matrices for accelerated attention benchmark results.""" + +from __future__ import annotations + +import argparse +import json +import math +from pathlib import Path + +import matplotlib.pyplot as plt + +_DEFAULT_INPUT = Path("artifacts/benchmark/flashdreams/accelerated/benchmark.json") +_DEFAULT_OUTPUT_DIR = Path("artifacts/benchmark/flashdreams/accelerated") + +RowKey = tuple[str, str, bool, bool] +CellKey = tuple[RowKey, str] + + +def _parse_args(argv: list[str] | None = None) -> argparse.Namespace: + """Parse benchmark input and plot output paths. + + Args: + argv: Command-line arguments; ``None`` reads ``sys.argv``. + + Returns: + Parsed command-line arguments. + """ + parser = argparse.ArgumentParser( + description="Plot self- and cross-attention median latency as PNG matrices." + ) + parser.add_argument( + "input", + nargs="?", + type=Path, + default=_DEFAULT_INPUT, + help=f"pytest-benchmark JSON path (default: {_DEFAULT_INPUT})", + ) + parser.add_argument( + "-o", + "--output-dir", + type=Path, + default=_DEFAULT_OUTPUT_DIR, + help=f"output directory (default: {_DEFAULT_OUTPUT_DIR})", + ) + return parser.parse_args(argv) + + +def _load_matrix( + input_path: Path, +) -> tuple[list[RowKey], list[str], dict[CellKey, float], str]: + """Load accelerated attention rows and median timings from benchmark JSON. + + Args: + input_path: Pytest-benchmark JSON file to parse. + + Returns: + Ordered row keys, implementation columns, median milliseconds by cell, + and plot subtitle. + + Raises: + SystemExit: The input cannot be read or does not contain compatible + accelerated attention records. + """ + try: + payload = json.loads(input_path.read_text(encoding="utf-8")) + except OSError as error: + raise SystemExit(f"Cannot read benchmark JSON {input_path}: {error}") from error + except json.JSONDecodeError as error: + raise SystemExit( + f"Cannot parse benchmark JSON {input_path}: {error}" + ) from error + + if not isinstance(payload, dict) or not isinstance(payload.get("benchmarks"), list): + raise SystemExit(f"Benchmark JSON {input_path} has no benchmarks list") + + rows: list[RowKey] = [] + columns: list[str] = [] + values_ms: dict[CellKey, float] = {} + first_extra: dict[str, object] | None = None + + for record in payload["benchmarks"]: + if not isinstance(record, dict): + continue + extra = record.get("extra_info") + if not isinstance(extra, dict) or not { + "attention_type", + "implementation_case", + "qk_norm_scope", + "rope_interleaved", + "bias", + }.issubset(extra): + continue + + attention_type = extra["attention_type"] + implementation = extra["implementation_case"] + norm = extra["qk_norm_scope"] + rope_interleaved = extra["rope_interleaved"] + bias = extra["bias"] + if not all( + isinstance(value, str) for value in (attention_type, implementation, norm) + ): + raise SystemExit("Attention labels in benchmark JSON must be strings") + if attention_type not in {"self_attention", "cross_attention"}: + raise SystemExit(f"Unsupported attention type {attention_type!r}") + if not isinstance(rope_interleaved, bool) or not isinstance(bias, bool): + raise SystemExit("Attention RoPE and bias settings must be booleans") + + stats = record.get("stats") + median = stats.get("median") if isinstance(stats, dict) else None + if ( + isinstance(median, bool) + or not isinstance(median, (int, float)) + or not math.isfinite(median) + or median <= 0 + ): + raise SystemExit( + f"Benchmark {record.get('name', '')} has no positive median" + ) + + row = (attention_type, norm, rope_interleaved, bias) + cell = (row, implementation) + if cell in values_ms: + raise SystemExit(f"Duplicate benchmark cell for {row} and {implementation}") + if row not in rows: + rows.append(row) + if implementation not in columns: + columns.append(implementation) + values_ms[cell] = median * 1000.0 + if first_extra is None: + first_extra = extra + + if not values_ms or first_extra is None: + raise SystemExit( + f"Benchmark JSON {input_path} contains no accelerated attention results" + ) + + subtitle_parts = ["Median latency in ms (lower is faster)"] + gpu = first_extra.get("gpu") + if isinstance(gpu, str) and gpu: + subtitle_parts.append(gpu) + commit_info = payload.get("commit_info") + commit_id = commit_info.get("id") if isinstance(commit_info, dict) else None + if isinstance(commit_id, str) and commit_id: + subtitle_parts.append(f"commit {commit_id[:10]}") + timestamp = payload.get("datetime") + if isinstance(timestamp, str) and timestamp: + subtitle_parts.append(timestamp) + return rows, columns, values_ms, " · ".join(subtitle_parts) + + +def _row_label(row: RowKey) -> str: + _, norm, rope_interleaved, bias = row + rope = "interleaved" if rope_interleaved else "split" + return f"norm {norm} | rope {rope} | bias {'on' if bias else 'off'}" + + +def _column_label(column: str) -> str: + return " ".join( + "fuse qkv" if token == "full" else token for token in column.split("-") + ) + + +def _cell_label( + value: float | None, reference: float | None, *, is_reference: bool +) -> str: + """Format a latency and its relationship to the row reference. + + Args: + value: Cell latency in milliseconds; ``None`` marks a missing result. + reference: First-column latency in milliseconds; ``None`` marks a + missing reference. + is_reference: Whether the cell is the first column in its row. + + Returns: + Two-line latency and relative-performance annotation. + """ + if value is None: + return "N/A" + if is_reference: + return f"{value:.2f} ms\n1.00× reference" + if reference is None: + return f"{value:.2f} ms\nreference unavailable" + if value < reference: + return f"{value:.2f} ms\n{(reference / value - 1) * 100:.0f}% faster" + if value > reference: + return f"{value:.2f} ms\n{(value / reference - 1) * 100:.0f}% slower" + return f"{value:.2f} ms\nsame as reference" + + +def _write_png( + output_path: Path, + attention_type: str, + rows: list[RowKey], + columns: list[str], + values_ms: dict[CellKey, float], + subtitle: str, +) -> None: + """Write one attention type's median-latency heatmap as a PNG. + + Args: + output_path: Destination PNG file. + attention_type: Attention family represented by every row. + rows: Ordered attention policy rows. + columns: Ordered implementation configuration columns. + values_ms: Median milliseconds keyed by row and implementation. + subtitle: Benchmark environment summary. + """ + attention_label = attention_type.replace("_", " ").title() + matrix = [ + [values_ms.get((row, column), math.nan) for column in columns] for row in rows + ] + figure, axes = plt.subplots() + default_width, default_height = figure.get_size_inches() + # ponytail: Linear sizing assumes current short labels; measure rendered + # text extents if benchmark labels become substantially longer. + figure.set_size_inches( + max(default_width, len(columns) * 2.0), + max(default_height, len(rows) * 0.75), + ) + image = axes.imshow(matrix, aspect="auto", cmap="Blues") + axes.set_xticks( + range(len(columns)), + labels=[_column_label(column) for column in columns], + rotation=45, + ha="right", + rotation_mode="anchor", + ) + axes.set_yticks(range(len(rows)), labels=[_row_label(row) for row in rows]) + axes.set_xlabel("Implementation configuration") + axes.set_ylabel("Attention policy") + axes.set_title(f"{attention_label} performance\n{subtitle}") + for row_index, row in enumerate(rows): + reference = values_ms.get((row, columns[0])) + for column_index, column in enumerate(columns): + value = values_ms.get((row, column)) + text_color = "black" + if value is not None: + red, green, blue, _ = image.cmap(image.norm(value)) + channels = (red, green, blue) + linear = tuple( + channel / 12.92 + if channel <= 0.04045 + else ((channel + 0.055) / 1.055) ** 2.4 + for channel in channels + ) + luminance = 0.2126 * linear[0] + 0.7152 * linear[1] + 0.0722 * linear[2] + text_color = "white" if luminance < 0.179 else "black" + axes.text( + column_index, + row_index, + _cell_label(value, reference, is_reference=column_index == 0), + ha="center", + va="center", + color=text_color, + ) + colorbar = figure.colorbar(image, ax=axes, label="Median latency (ms)") + colorbar.ax.text(0.5, 1.02, "Slower ↑", ha="center", transform=colorbar.ax.transAxes) + colorbar.ax.text( + 0.5, + -0.04, + "↓ Faster", + ha="center", + va="top", + transform=colorbar.ax.transAxes, + ) + figure.tight_layout() + output_path.parent.mkdir(parents=True, exist_ok=True) + figure.savefig(output_path) + plt.close(figure) + + +def main(argv: list[str] | None = None) -> None: + """Generate separate self- and cross-attention performance matrices.""" + args = _parse_args(argv) + rows, columns, values_ms, subtitle = _load_matrix(args.input) + attention_types = list(dict.fromkeys(row[0] for row in rows)) + for attention_type in attention_types: + panel_rows = [row for row in rows if row[0] == attention_type] + panel_values = { + cell: value + for cell, value in values_ms.items() + if cell[0][0] == attention_type + } + output = args.output_dir / f"{attention_type}.png" + _write_png( + output, + attention_type, + panel_rows, + columns, + panel_values, + subtitle, + ) + print(f"Wrote {output}") + + +if __name__ == "__main__": + main() diff --git a/scripts/benchmark/flashdreams/accelerated/run.sh b/scripts/benchmark/flashdreams/accelerated/run.sh new file mode 100755 index 000000000..4edb4101a --- /dev/null +++ b/scripts/benchmark/flashdreams/accelerated/run.sh @@ -0,0 +1,13 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +set -euo pipefail + +cd "$(cd "$(dirname "${BASH_SOURCE[0]}")/../../../.." && pwd)" +mkdir -p artifacts/benchmark/flashdreams/accelerated + +uv run --project flashdreams --group test pytest \ + flashdreams/benchmarks/accelerated \ + -p no:manual_marker -m manual --benchmark-only -v "$@" \ + --benchmark-json=artifacts/benchmark/flashdreams/accelerated/benchmark.json diff --git a/scripts/benchmark/flashdreams/recipes/run.sh b/scripts/benchmark/flashdreams/recipes/run.sh new file mode 100755 index 000000000..d0f80f2b8 --- /dev/null +++ b/scripts/benchmark/flashdreams/recipes/run.sh @@ -0,0 +1,12 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +set -euo pipefail + +cd "$(cd "$(dirname "${BASH_SOURCE[0]}")/../../../.." && pwd)" +mkdir -p artifacts/benchmark/flashdreams/recipes + +uv run --project flashdreams --group test pytest flashdreams/benchmarks/recipes \ + -p no:manual_marker -m manual --benchmark-only -v "$@" \ + --benchmark-json=artifacts/benchmark/flashdreams/recipes/benchmark.json diff --git a/scripts/benchmark/lingbot/run.sh b/scripts/benchmark/lingbot/run.sh new file mode 100755 index 000000000..4f2d9e58a --- /dev/null +++ b/scripts/benchmark/lingbot/run.sh @@ -0,0 +1,13 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +set -euo pipefail + +cd "$(cd "$(dirname "${BASH_SOURCE[0]}")/../../.." && pwd)" +mkdir -p artifacts/benchmark/lingbot + +uv run --project integrations/lingbot --group test pytest \ + integrations/lingbot/benchmarks \ + -p no:manual_marker -m manual --benchmark-only -v "$@" \ + --benchmark-json=artifacts/benchmark/lingbot/benchmark.json diff --git a/scripts/benchmark/omnidreams/plot.py b/scripts/benchmark/omnidreams/plot.py new file mode 100644 index 000000000..ed5340287 --- /dev/null +++ b/scripts/benchmark/omnidreams/plot.py @@ -0,0 +1,392 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Bar plots for Omnidreams module and end-to-end benchmark results.""" + +from __future__ import annotations + +import argparse +import json +import math +from pathlib import Path + +import matplotlib.pyplot as plt + +_DEFAULT_INPUT = Path("artifacts/benchmark/omnidreams/benchmark.json") +_DEFAULT_OUTPUT_DIR = Path("artifacts/benchmark/omnidreams") + +_MODULE_PANELS = ( + ("omnidreams-dit-self-attention", "Self-attention"), + ("omnidreams-dit-cross-attention", "Cross-attention"), + ("omnidreams-dit-block", "DiT block"), +) +_END_TO_END_PANELS = ( + ("omnidreams-dit-network", "Network eval"), + ("omnidreams-full-pipeline-generate", "Pipeline generate"), + ("omnidreams-full-pipeline-finalize", "Pipeline finalize"), +) +_PANELS = (*_MODULE_PANELS, *_END_TO_END_PANELS) + +_END_TO_END_LABELS = { + "omnidreams_torch": "Omnidreams PyTorch\nself + cross", + "triton_fa2_fp8_full": "Triton FA2 FP8\nQKV self + KV cross", + "triton_cudnn_bf16_full": "Triton cuDNN BF16\nQKV self + KV cross", + "triton_cudnn_bf16_full_omnidreams_cross": ( + "Triton cuDNN BF16\nQKV self + Omnidreams cross" + ), + "cuda": "Native CUDA FP8\ncuDNN", + "cuda_sparge": "Native CUDA FP8\nSparge", + "cuda_sage3": "Native CUDA BF16\nSage3", + "cuda_sage3_fp8": "Native CUDA FP8\nSage3", +} +"""Compact labels for the selected mixed self/cross end-to-end configurations.""" + +BenchmarkValues = dict[str, dict[str, float]] +Panel = tuple[str, str] + + +def _parse_args(argv: list[str] | None = None) -> argparse.Namespace: + """Parse benchmark input and plot output paths. + + Args: + argv: Command-line arguments; ``None`` reads ``sys.argv``. + + Returns: + Parsed command-line arguments. + """ + parser = argparse.ArgumentParser( + description="Plot Omnidreams median benchmark latency as bar charts." + ) + parser.add_argument( + "input", + nargs="?", + type=Path, + default=_DEFAULT_INPUT, + help=f"pytest-benchmark JSON path (default: {_DEFAULT_INPUT})", + ) + parser.add_argument( + "-o", + "--output-dir", + type=Path, + default=_DEFAULT_OUTPUT_DIR, + help=f"output directory (default: {_DEFAULT_OUTPUT_DIR})", + ) + return parser.parse_args(argv) + + +def _load_results(input_path: Path) -> tuple[BenchmarkValues, list[str], str]: + """Load Omnidreams median timings and environment metadata. + + Args: + input_path: Pytest-benchmark JSON file to parse. + + Returns: + Median milliseconds by group and implementation, configuration order, + and plot subtitle. + + Raises: + SystemExit: The input cannot be read or lacks complete Omnidreams data. + """ + try: + payload = json.loads(input_path.read_text(encoding="utf-8")) + except OSError as error: + raise SystemExit(f"Cannot read benchmark JSON {input_path}: {error}") from error + except json.JSONDecodeError as error: + raise SystemExit( + f"Cannot parse benchmark JSON {input_path}: {error}" + ) from error + + if not isinstance(payload, dict) or not isinstance(payload.get("benchmarks"), list): + raise SystemExit(f"Benchmark JSON {input_path} has no benchmarks list") + + values: BenchmarkValues = {group: {} for group, _ in _PANELS} + configurations: list[str] = [] + first_extra: dict[str, object] | None = None + + for record in payload["benchmarks"]: + if not isinstance(record, dict): + continue + group = record.get("group") + if not isinstance(group, str) or group not in values: + continue + extra = record.get("extra_info") + implementation = ( + extra.get("implementation") if isinstance(extra, dict) else None + ) + if not isinstance(implementation, str) or not implementation: + raise SystemExit( + f"Benchmark {record.get('name', '')} has no implementation" + ) + + stats = record.get("stats") + median = stats.get("median") if isinstance(stats, dict) else None + if ( + isinstance(median, bool) + or not isinstance(median, (int, float)) + or not math.isfinite(median) + or median <= 0 + ): + raise SystemExit( + f"Benchmark {record.get('name', '')} has no positive median" + ) + if implementation in values[group]: + raise SystemExit(f"Duplicate benchmark for {group} and {implementation}") + + values[group][implementation] = median * 1000.0 + if implementation not in configurations: + configurations.append(implementation) + if first_extra is None: + first_extra = extra + + missing = [group for group, results in values.items() if not results] + if missing or first_extra is None: + raise SystemExit( + f"Benchmark JSON {input_path} lacks groups: {', '.join(missing)}" + ) + return values, configurations, _subtitle(payload, first_extra) + + +def _subtitle(payload: dict[str, object], extra: dict[str, object]) -> str: + """Build a compact benchmark-environment subtitle.""" + parts = ["Median latency in ms (lower is faster)"] + gpu = extra.get("gpu") + if isinstance(gpu, str) and gpu: + parts.append(gpu) + warmups = extra.get("warmup_rounds") + rounds = extra.get("benchmark_rounds") + if isinstance(warmups, int) and isinstance(rounds, int): + parts.append(f"{warmups} warmups / {rounds} rounds") + commit_info = payload.get("commit_info") + commit_id = commit_info.get("id") if isinstance(commit_info, dict) else None + if isinstance(commit_id, str) and commit_id: + parts.append(f"commit {commit_id[:10]}") + timestamp = payload.get("datetime") + if isinstance(timestamp, str) and timestamp: + parts.append(timestamp) + return " · ".join(parts) + + +def _configuration_label( + implementation: str, + *, + end_to_end: bool = False, +) -> str: + """Format a stable implementation name as a compact axis label. + + Args: + implementation: Stable benchmark implementation identifier. + end_to_end: Whether to describe both self- and cross-attention branches. + + Returns: + Compact multi-line label for the plot axis. + """ + if end_to_end: + label = _END_TO_END_LABELS.get(implementation) + if label is not None: + return label + if implementation == "triton_cudnn_bf16_full_omnidreams_cross": + return _END_TO_END_LABELS[implementation] + if implementation == "omnidreams_torch": + return "Omnidreams\nPyTorch" + if implementation == "cuda": + return "Native\nCUDA" + if implementation.startswith("triton_"): + parts = implementation.removeprefix("triton_").split("_", maxsplit=2) + if len(parts) == 3: + backend, precision, fusion = parts + backend_label = "cuDNN" if backend == "cudnn" else backend.upper() + fusion_label = ( + "FUSE QKV" if fusion == "full" else fusion.replace("_", " ").upper() + ) + return f"{backend_label}\n{precision.upper()}\n{fusion_label}" + return implementation.replace("_", "\n") + + +def _bar_label( + latency_ms: float, + reference_ms: float, + *, + is_reference: bool, +) -> str: + """Format latency and relative performance against PyTorch. + + Args: + latency_ms: Bar latency in milliseconds. + reference_ms: PyTorch reference latency in milliseconds. + is_reference: Whether the bar is the PyTorch reference. + + Returns: + Two-line latency and relative-performance annotation. + """ + if latency_ms < 1: + latency_label = f"{latency_ms:.3f} ms" + elif latency_ms < 10: + latency_label = f"{latency_ms:.2f} ms" + else: + latency_label = f"{latency_ms:.1f} ms" + if is_reference: + return f"{latency_label}\nreference" + if latency_ms < reference_ms: + return f"{latency_label}\n{(reference_ms / latency_ms - 1) * 100:.0f}% faster" + if latency_ms > reference_ms: + return f"{latency_label}\n{(latency_ms / reference_ms - 1) * 100:.0f}% slower" + return f"{latency_label}\nsame as reference" + + +def _panel_configurations( + panels: tuple[Panel, ...], + values: BenchmarkValues, + configurations: list[str], +) -> list[str]: + return [ + configuration + for configuration in configurations + if any(configuration in values[group] for group, _ in panels) + ] + + +def _write_bar_figure( + output_path: Path, + title: str, + panels: tuple[Panel, ...], + values: BenchmarkValues, + configurations: list[str], + subtitle: str, +) -> None: + """Write three aligned median-latency bar charts as one PNG. + + Args: + output_path: Destination PNG file. + title: Figure title. + panels: Benchmark group and display-title pairs. + values: Median milliseconds by group and implementation. + configurations: Stable implementation order. + subtitle: Benchmark environment summary. + """ + end_to_end = panels == _END_TO_END_PANELS + panel_configurations = _panel_configurations(panels, values, configurations) + figure, axes = plt.subplots( + len(panels), + 1, + figsize=(max(16.0, len(panel_configurations) * 1.3), 15.0), + layout="constrained", + ) + axes_list = list(axes) + for axes_item, (group, panel_title) in zip(axes_list, panels, strict=True): + reference = values[group]["omnidreams_torch"] + ordered_configurations = sorted( + panel_configurations, + key=lambda configuration: values[group].get(configuration, math.inf), + ) + latencies = [ + values[group].get(configuration, math.nan) + for configuration in ordered_configurations + ] + colors = [ + "C2" + if configuration == "omnidreams_torch" + else "C1" + if configuration == "cuda" + else "C0" + for configuration in ordered_configurations + ] + bars = axes_item.bar( + range(len(ordered_configurations)), + latencies, + color=colors, + ) + finite_latencies = [value for value in latencies if math.isfinite(value)] + axes_item.set_ylim(0, max(finite_latencies) * 1.2) + axes_item.set_title(panel_title) + axes_item.set_ylabel("Median latency (ms)") + + for index, (bar, latency, configuration) in enumerate( + zip(bars, latencies, ordered_configurations, strict=True) + ): + if not math.isfinite(latency): + bar.set_visible(False) + axes_item.text( + index, + 0.02, + "N/A", + transform=axes_item.get_xaxis_transform(), + ha="center", + va="bottom", + ) + continue + axes_item.annotate( + _bar_label( + latency, + reference, + is_reference=configuration == "omnidreams_torch", + ), + xy=(bar.get_x() + bar.get_width() / 2, latency), + xytext=(0, 3), + textcoords="offset points", + ha="center", + va="bottom", + ) + axes_item.set_xticks( + range(len(ordered_configurations)), + labels=[ + _configuration_label( + configuration, + end_to_end=end_to_end, + ) + for configuration in ordered_configurations + ], + ) + + figure.supxlabel( + "Configuration (self-attention + cross-attention)" + if end_to_end + else "Configuration (SDPA backend / precision / QKV fusion)" + ) + figure.suptitle(f"{title}\n{subtitle}") + output_path.parent.mkdir(parents=True, exist_ok=True) + figure.savefig(output_path) + plt.close(figure) + + +def main(argv: list[str] | None = None) -> None: + """Generate module and end-to-end Omnidreams benchmark figures.""" + args = _parse_args(argv) + values, configurations, subtitle = _load_results(args.input) + outputs = ( + ( + args.output_dir / "modules.png", + "Omnidreams module benchmarks", + _MODULE_PANELS, + ), + ( + args.output_dir / "network_pipeline.png", + "Omnidreams network and pipeline benchmarks", + _END_TO_END_PANELS, + ), + ) + for output_path, title, panels in outputs: + _write_bar_figure( + output_path, + title, + panels, + values, + configurations, + subtitle, + ) + print(f"Wrote {output_path}") + + +if __name__ == "__main__": + main() diff --git a/scripts/benchmark/omnidreams/run.sh b/scripts/benchmark/omnidreams/run.sh new file mode 100755 index 000000000..67a5a0b8b --- /dev/null +++ b/scripts/benchmark/omnidreams/run.sh @@ -0,0 +1,13 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +set -euo pipefail + +cd "$(cd "$(dirname "${BASH_SOURCE[0]}")/../../.." && pwd)" +mkdir -p artifacts/benchmark/omnidreams + +uv run --project integrations/omnidreams --group test pytest \ + integrations/omnidreams/benchmarks \ + -p no:manual_marker -m manual --benchmark-only -v "$@" \ + --benchmark-json=artifacts/benchmark/omnidreams/benchmark.json diff --git a/scripts/benchmark/run_all.sh b/scripts/benchmark/run_all.sh new file mode 100755 index 000000000..2065b2050 --- /dev/null +++ b/scripts/benchmark/run_all.sh @@ -0,0 +1,15 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +set -euo pipefail +shopt -s globstar nullglob + +script_dir=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) + +status=0 +for run_script in "$script_dir"/**/run.sh; do + "$run_script" "$@" || status=$? +done + +exit "$status" diff --git a/scripts/benchmark/wan21/run.sh b/scripts/benchmark/wan21/run.sh new file mode 100755 index 000000000..459a1985a --- /dev/null +++ b/scripts/benchmark/wan21/run.sh @@ -0,0 +1,13 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +set -euo pipefail + +cd "$(cd "$(dirname "${BASH_SOURCE[0]}")/../../.." && pwd)" +mkdir -p artifacts/benchmark/wan21 + +uv run --project integrations/wan21 --group test pytest \ + integrations/wan21/benchmarks \ + -p no:manual_marker -m manual --benchmark-only -v "$@" \ + --benchmark-json=artifacts/benchmark/wan21/benchmark.json diff --git a/uv.lock b/uv.lock index d894ab6a8..a48172c2d 100644 --- a/uv.lock +++ b/uv.lock @@ -42,6 +42,7 @@ overrides = [ ] [manifest.dependency-groups] +dev = [{ name = "python-dotenv", specifier = ">=1.2.2" }] docs = [ { name = "myst-parser", specifier = ">=4.0" }, { name = "pydata-sphinx-theme", specifier = ">=0.18" }, @@ -65,7 +66,9 @@ lint = [ { name = "pre-commit", specifier = ">=4.3.0" }, { name = "pytest", specifier = ">=8.0" }, { name = "pytest-asyncio", specifier = ">=0.23" }, + { name = "pytest-benchmark", specifier = ">=5.1" }, { name = "pytest-manual-marker", specifier = ">=2.0" }, + { name = "ruff", specifier = "==0.12.7" }, { name = "sphinx", specifier = ">=7.0" }, { name = "tomli", specifier = ">=2.0" }, { name = "ty", specifier = ">=0.0.39" }, @@ -74,6 +77,7 @@ test = [ { name = "imageio-ffmpeg", specifier = ">=0.5" }, { name = "pytest", specifier = ">=8.0" }, { name = "pytest-asyncio", specifier = ">=0.23" }, + { name = "pytest-benchmark", specifier = ">=5.1" }, { name = "pytest-manual-marker", specifier = ">=2.0" }, { name = "tomli", specifier = ">=2.0" }, ] @@ -998,6 +1002,7 @@ dependencies = [ { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or (extra == 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12') or (extra == 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13')" }, { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' or (extra == 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12') or (extra == 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13')" }, { name = "nvidia-ml-py" }, + { name = "nvtx" }, { name = "psutil" }, { name = "pyyaml" }, { name = "safetensors" }, @@ -1075,6 +1080,7 @@ requires-dist = [ { name = "numpy", specifier = ">=1.24,<2.5" }, { name = "nvidia-ml-py", specifier = ">=12.0" }, { name = "nvidia-vfx", marker = "extra == 'rtx-postprocess'", specifier = "==0.1.0.1" }, + { name = "nvtx", specifier = ">=0.2.15" }, { name = "opencv-python-headless", marker = "extra == 'examples'", specifier = ">=4.5" }, { name = "opencv-python-headless", marker = "extra == 'runners'", specifier = ">=4.5" }, { name = "psutil", specifier = ">=7.0" }, @@ -3093,6 +3099,29 @@ version = "0.1.0.1" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/37/b4/58e1bbb8d6fc9ed786564d0878314ea2f2cd458c84861a5130927e431ff6/nvidia_vfx-0.1.0.1.tar.gz", hash = "sha256:8a26bae3a967a2ce29040f17ba9d75e106f3d0c68016d440a77ed9c7eb05daae", size = 2673, upload-time = "2026-03-09T19:29:40.556Z" } +[[package]] +name = "nvtx" +version = "0.2.15" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/92/dd/692765e87de30bae1522cdffaa0f2b52949658a92a0fa6d96b1a01eae9d2/nvtx-0.2.15.tar.gz", hash = "sha256:2287d3be05b85661deb386f878d1f536c2e532774aa9ec7a50c434942ed81ae5", size = 121230, upload-time = "2026-03-18T10:01:25.547Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9d/ef/ea1e9d92afd07fdf2a2390e508f1d214e5ba890561d7849d6ca708534b9d/nvtx-0.2.15-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a4f50832fd90a1b480a9deef6e4cd48015b61869095b54dd1a7afe87b4138c6a", size = 768543, upload-time = "2026-03-18T10:07:21.819Z" }, + { url = "https://files.pythonhosted.org/packages/32/8e/b42c05cf3cc43c51f21fdda6f7c4fe28a595c6d2bdb0cfbf0477dc5805f2/nvtx-0.2.15-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:5f3362f0db4252514719326c9d5662b0f93d254659ba97b9c8dbe556286e0e3e", size = 771975, upload-time = "2026-03-18T10:12:23.772Z" }, + { url = "https://files.pythonhosted.org/packages/60/77/fc000055b5bb1651cdd772f0fe1fd9a16c7773b28dfc5624eea331d1415d/nvtx-0.2.15-cp310-cp310-win_amd64.whl", hash = "sha256:d71f934e580d4572f382712b6da464ab69e4c212981506f781f927d5c6d935d6", size = 134503, upload-time = "2026-03-18T10:04:05.773Z" }, + { url = "https://files.pythonhosted.org/packages/80/65/435d10b2041ee082c07d5aed129afd504012c8908796d695f10e66bcc716/nvtx-0.2.15-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:157b80ea9b4db6c8f47f8dbe2fa2e81e7a7f1445bb87f8268f43dec9210b78a1", size = 806443, upload-time = "2026-03-18T10:05:49.308Z" }, + { url = "https://files.pythonhosted.org/packages/47/bc/be94576ba33af75bcc68a857daade64cb86481764d4fb0f36308b1f6fc85/nvtx-0.2.15-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:02bca69ee55e0be41eabf908de9dbcdd18e702c7f49f9aa63fd396ce684ff5d5", size = 808183, upload-time = "2026-03-18T10:11:16.262Z" }, + { url = "https://files.pythonhosted.org/packages/f6/7a/42109f1cfb1ff9913201cb2b804956a4f003db4c018c2522a3c8066b3a1c/nvtx-0.2.15-cp311-cp311-win_amd64.whl", hash = "sha256:dbe41f78f5a811bd4cdad0a237e5b41a4937d8c2c6c9abdd161091671a598bc0", size = 134631, upload-time = "2026-03-18T10:02:11.247Z" }, + { url = "https://files.pythonhosted.org/packages/c2/07/698355285a03a366ef63ea9762fc1feef3f9f25483e1655408f72d827090/nvtx-0.2.15-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:2cc530cd0f1a2c14a3a7e683833db509888ac5ed4ead94e5c9e2c7317c6937a7", size = 807159, upload-time = "2026-03-18T10:09:49.232Z" }, + { url = "https://files.pythonhosted.org/packages/c0/d1/08f22448d83481408d663065764ba583df091a7de629ed38fc97e522f1af/nvtx-0.2.15-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:3ca8030a6d197952318013dd1c12c22da1d4b9feb76ba72e0fcd449961183c2c", size = 806187, upload-time = "2026-03-18T10:13:32.972Z" }, + { url = "https://files.pythonhosted.org/packages/54/23/c97c39e3b7ba256aa343cb828ca0d1c8421f705ca84795658ecd14ca95ed/nvtx-0.2.15-cp312-cp312-win_amd64.whl", hash = "sha256:70a1e768964e0520b68ccabc4df391cc227537c45936a7eba6507bc65e617e00", size = 129178, upload-time = "2026-03-18T10:02:55.299Z" }, + { url = "https://files.pythonhosted.org/packages/05/c9/8341224b8284f7deb6a634119939de5885adc421e64b6743693b30da2186/nvtx-0.2.15-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d28660d9c46f8ba750d781572b6aa5a1e6221abba224ab32d7fb32c2d0fd67df", size = 780787, upload-time = "2026-03-18T10:10:40.634Z" }, + { url = "https://files.pythonhosted.org/packages/b1/c0/4a5bb7897918de7c7e0191d9342df8ae4cb797ff07276e0f20d13e497ce7/nvtx-0.2.15-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:10749686633f880ad53dcdbb2179fad41b45dcf5b7631d4a1070a577577bd386", size = 782575, upload-time = "2026-03-18T10:13:57.3Z" }, + { url = "https://files.pythonhosted.org/packages/38/b9/6b381ac7c5a3ded331aebbf25f8959d19b51d320fb2514c76c6b6edddaaa/nvtx-0.2.15-cp313-cp313-win_amd64.whl", hash = "sha256:a6650b029263d12f8427a4dee8bd59cb9c91bccb60543bfcb20bc2b00fdcd672", size = 128764, upload-time = "2026-03-18T10:02:33.343Z" }, + { url = "https://files.pythonhosted.org/packages/75/69/a9acb6d95d2e0e381b2956544768528dd8d7a9e827af8c2014169d838284/nvtx-0.2.15-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:25813ead4fff4d3a6e04f69a72507b096a6bdbecefa369f1100b0e584767bca8", size = 833375, upload-time = "2026-03-18T10:06:31.955Z" }, + { url = "https://files.pythonhosted.org/packages/38/56/c7e8645061cc2fc23f3a54f33e1e340df59216f07dcfb97d46b8ae7dd26c/nvtx-0.2.15-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:3741edac4678b92f03d22a3f0a2dfd469f422f85e63db71b038e02525b2404ad", size = 788639, upload-time = "2026-03-18T10:12:01.69Z" }, + { url = "https://files.pythonhosted.org/packages/96/03/fadd82acdbca6d1c49ac517081a0c3714346f52f4c7e1d4449d77605b4aa/nvtx-0.2.15-cp313-cp313t-win_amd64.whl", hash = "sha256:8be06c3c8c267eba56a0396366b9593092e0b75ea8d3702b303d48c0a1662f0e", size = 142609, upload-time = "2026-03-18T10:01:48.832Z" }, +] + [[package]] name = "onnx" version = "1.22.0" @@ -3499,6 +3528,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/8e/37/efad0257dc6e593a18957422533ff0f87ede7c9c6ea010a2177d738fb82f/pure_eval-0.2.3-py3-none-any.whl", hash = "sha256:1db8e35b67b3d218d818ae653e27f06c3aa420901fa7b081ca98cbedc874e0d0", size = 11842, upload-time = "2024-07-21T12:58:20.04Z" }, ] +[[package]] +name = "py-cpuinfo" +version = "9.0.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/37/a8/d832f7293ebb21690860d2e01d8115e5ff6f2ae8bbdc953f0eb0fa4bd2c7/py-cpuinfo-9.0.0.tar.gz", hash = "sha256:3cdbbf3fac90dc6f118bfd64384f309edeadd902d7c8fb17f02ffa1fc3f49690", size = 104716, upload-time = "2022-10-25T20:38:06.303Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e0/a9/023730ba63db1e494a271cb018dcd361bd2c917ba7004c3e49d5daf795a2/py_cpuinfo-9.0.0-py3-none-any.whl", hash = "sha256:859625bc251f64e21f077d099d4162689c762b5d6a4c3c97553d56241c9674d5", size = 22335, upload-time = "2022-10-25T20:38:27.636Z" }, +] + [[package]] name = "pyarrow" version = "24.0.0" @@ -3757,6 +3795,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/03/e2/08a497ef684b88559c9cc5f4ad53a37e7b99e727094a86d6ea32536d5d3c/pytest_asyncio-1.4.0-py3-none-any.whl", hash = "sha256:933ca923a23075a87fb7070c0ec272a6848489824d887c85c812670932835aa1", size = 16930, upload-time = "2026-05-26T09:56:02.576Z" }, ] +[[package]] +name = "pytest-benchmark" +version = "5.2.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "py-cpuinfo" }, + { name = "pytest" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/24/34/9f732b76456d64faffbef6232f1f9dbec7a7c4999ff46282fa418bd1af66/pytest_benchmark-5.2.3.tar.gz", hash = "sha256:deb7317998a23c650fd4ff76e1230066a76cb45dcece0aca5607143c619e7779", size = 341340, upload-time = "2025-11-09T18:48:43.215Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/33/29/e756e715a48959f1c0045342088d7ca9762a2f509b945f362a316e9412b7/pytest_benchmark-5.2.3-py3-none-any.whl", hash = "sha256:bc839726ad20e99aaa0d11a127445457b4219bdb9e80a1afc4b51da7f96b0803", size = 45255, upload-time = "2025-11-09T18:48:39.765Z" }, +] + [[package]] name = "pytest-manual-marker" version = "2.0.0.0" @@ -3795,6 +3846,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/1a/82/a70006589557f267f15bd384c0642ad49f0d97b690c3a05b166b9dcbad3b/python_discovery-1.4.2-py3-none-any.whl", hash = "sha256:475803f53b7b2ed6e490e27373f9d8340f7d2eebf9acdaf645d7d714c97bb500", size = 33886, upload-time = "2026-06-11T16:10:41.192Z" }, ] +[[package]] +name = "python-dotenv" +version = "1.2.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/82/ed/0301aeeac3e5353ef3d94b6ec08bbcabd04a72018415dcb29e588514bba8/python_dotenv-1.2.2.tar.gz", hash = "sha256:2c371a91fbd7ba082c2c1dc1f8bf89ca22564a087c2c287cd9b662adde799cf3", size = 50135, upload-time = "2026-03-01T16:00:26.196Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0b/d7/1959b9648791274998a9c3526f6d0ec8fd2233e4d4acce81bbae76b44b2a/python_dotenv-1.2.2-py3-none-any.whl", hash = "sha256:1d8214789a24de455a8b8bd8ae6fe3c6b69a5e3d64aa8a8e5d68e694bbcb285a", size = 22101, upload-time = "2026-03-01T16:00:25.09Z" }, +] + [[package]] name = "pytz" version = "2026.2" @@ -3957,6 +4017,31 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/04/54/6f679c435d28e0a568d8e8a7c0a93a09010818634c3c3907fc98d8983770/roman_numerals-4.1.0-py3-none-any.whl", hash = "sha256:647ba99caddc2cc1e55a51e4360689115551bf4476d90e8162cf8c345fe233c7", size = 7676, upload-time = "2025-12-17T18:25:33.098Z" }, ] +[[package]] +name = "ruff" +version = "0.12.7" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/a1/81/0bd3594fa0f690466e41bd033bdcdf86cba8288345ac77ad4afbe5ec743a/ruff-0.12.7.tar.gz", hash = "sha256:1fc3193f238bc2d7968772c82831a4ff69252f673be371fb49663f0068b7ec71", size = 5197814, upload-time = "2025-07-29T22:32:35.877Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e1/d2/6cb35e9c85e7a91e8d22ab32ae07ac39cc34a71f1009a6f9e4a2a019e602/ruff-0.12.7-py3-none-linux_armv6l.whl", hash = "sha256:76e4f31529899b8c434c3c1dede98c4483b89590e15fb49f2d46183801565303", size = 11852189, upload-time = "2025-07-29T22:31:41.281Z" }, + { url = "https://files.pythonhosted.org/packages/63/5b/a4136b9921aa84638f1a6be7fb086f8cad0fde538ba76bda3682f2599a2f/ruff-0.12.7-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:789b7a03e72507c54fb3ba6209e4bb36517b90f1a3569ea17084e3fd295500fb", size = 12519389, upload-time = "2025-07-29T22:31:54.265Z" }, + { url = "https://files.pythonhosted.org/packages/a8/c9/3e24a8472484269b6b1821794141f879c54645a111ded4b6f58f9ab0705f/ruff-0.12.7-py3-none-macosx_11_0_arm64.whl", hash = "sha256:2e1c2a3b8626339bb6369116e7030a4cf194ea48f49b64bb505732a7fce4f4e3", size = 11743384, upload-time = "2025-07-29T22:31:59.575Z" }, + { url = "https://files.pythonhosted.org/packages/26/7c/458dd25deeb3452c43eaee853c0b17a1e84169f8021a26d500ead77964fd/ruff-0.12.7-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:32dec41817623d388e645612ec70d5757a6d9c035f3744a52c7b195a57e03860", size = 11943759, upload-time = "2025-07-29T22:32:01.95Z" }, + { url = "https://files.pythonhosted.org/packages/7f/8b/658798472ef260ca050e400ab96ef7e85c366c39cf3dfbef4d0a46a528b6/ruff-0.12.7-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:47ef751f722053a5df5fa48d412dbb54d41ab9b17875c6840a58ec63ff0c247c", size = 11654028, upload-time = "2025-07-29T22:32:04.367Z" }, + { url = "https://files.pythonhosted.org/packages/a8/86/9c2336f13b2a3326d06d39178fd3448dcc7025f82514d1b15816fe42bfe8/ruff-0.12.7-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:a828a5fc25a3efd3e1ff7b241fd392686c9386f20e5ac90aa9234a5faa12c423", size = 13225209, upload-time = "2025-07-29T22:32:06.952Z" }, + { url = "https://files.pythonhosted.org/packages/76/69/df73f65f53d6c463b19b6b312fd2391dc36425d926ec237a7ed028a90fc1/ruff-0.12.7-py3-none-manylinux_2_17_ppc64.manylinux2014_ppc64.whl", hash = "sha256:5726f59b171111fa6a69d82aef48f00b56598b03a22f0f4170664ff4d8298efb", size = 14182353, upload-time = "2025-07-29T22:32:10.053Z" }, + { url = "https://files.pythonhosted.org/packages/58/1e/de6cda406d99fea84b66811c189b5ea139814b98125b052424b55d28a41c/ruff-0.12.7-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:74e6f5c04c4dd4aba223f4fe6e7104f79e0eebf7d307e4f9b18c18362124bccd", size = 13631555, upload-time = "2025-07-29T22:32:12.644Z" }, + { url = "https://files.pythonhosted.org/packages/6f/ae/625d46d5164a6cc9261945a5e89df24457dc8262539ace3ac36c40f0b51e/ruff-0.12.7-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:5d0bfe4e77fba61bf2ccadf8cf005d6133e3ce08793bbe870dd1c734f2699a3e", size = 12667556, upload-time = "2025-07-29T22:32:15.312Z" }, + { url = "https://files.pythonhosted.org/packages/55/bf/9cb1ea5e3066779e42ade8d0cd3d3b0582a5720a814ae1586f85014656b6/ruff-0.12.7-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:06bfb01e1623bf7f59ea749a841da56f8f653d641bfd046edee32ede7ff6c606", size = 12939784, upload-time = "2025-07-29T22:32:17.69Z" }, + { url = "https://files.pythonhosted.org/packages/55/7f/7ead2663be5627c04be83754c4f3096603bf5e99ed856c7cd29618c691bd/ruff-0.12.7-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:e41df94a957d50083fd09b916d6e89e497246698c3f3d5c681c8b3e7b9bb4ac8", size = 11771356, upload-time = "2025-07-29T22:32:20.134Z" }, + { url = "https://files.pythonhosted.org/packages/17/40/a95352ea16edf78cd3a938085dccc55df692a4d8ba1b3af7accbe2c806b0/ruff-0.12.7-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:4000623300563c709458d0ce170c3d0d788c23a058912f28bbadc6f905d67afa", size = 11612124, upload-time = "2025-07-29T22:32:22.645Z" }, + { url = "https://files.pythonhosted.org/packages/4d/74/633b04871c669e23b8917877e812376827c06df866e1677f15abfadc95cb/ruff-0.12.7-py3-none-musllinux_1_2_i686.whl", hash = "sha256:69ffe0e5f9b2cf2b8e289a3f8945b402a1b19eff24ec389f45f23c42a3dd6fb5", size = 12479945, upload-time = "2025-07-29T22:32:24.765Z" }, + { url = "https://files.pythonhosted.org/packages/be/34/c3ef2d7799c9778b835a76189c6f53c179d3bdebc8c65288c29032e03613/ruff-0.12.7-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:a07a5c8ffa2611a52732bdc67bf88e243abd84fe2d7f6daef3826b59abbfeda4", size = 12998677, upload-time = "2025-07-29T22:32:27.022Z" }, + { url = "https://files.pythonhosted.org/packages/77/ab/aca2e756ad7b09b3d662a41773f3edcbd262872a4fc81f920dc1ffa44541/ruff-0.12.7-py3-none-win32.whl", hash = "sha256:c928f1b2ec59fb77dfdf70e0419408898b63998789cc98197e15f560b9e77f77", size = 11756687, upload-time = "2025-07-29T22:32:29.381Z" }, + { url = "https://files.pythonhosted.org/packages/b4/71/26d45a5042bc71db22ddd8252ca9d01e9ca454f230e2996bb04f16d72799/ruff-0.12.7-py3-none-win_amd64.whl", hash = "sha256:9c18f3d707ee9edf89da76131956aba1270c6348bfee8f6c647de841eac7194f", size = 12912365, upload-time = "2025-07-29T22:32:31.517Z" }, + { url = "https://files.pythonhosted.org/packages/4c/9b/0b8aa09817b63e78d94b4977f18b1fcaead3165a5ee49251c5d5c245bb2d/ruff-0.12.7-py3-none-win_arm64.whl", hash = "sha256:dfce05101dbd11833a0776716d5d1578641b7fddb537fe7fa956ab85d1769b69", size = 11982083, upload-time = "2025-07-29T22:32:33.881Z" }, +] + [[package]] name = "s3transfer" version = "0.19.0"