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/flashdreams/benchmarks/accelerated/multi_head_attention/test_multi_head_attention_benchmark.py b/flashdreams/benchmarks/accelerated/multi_head_attention/test_multi_head_attention_benchmark.py new file mode 100644 index 000000000..4151672d1 --- /dev/null +++ b/flashdreams/benchmarks/accelerated/multi_head_attention/test_multi_head_attention_benchmark.py @@ -0,0 +1,713 @@ +# 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. + +"""Matched microbenchmarks for Torch and optimized attention multi-head attention. + +Both implementations use identical geometry and random weights. + +Run the manual GPU benchmarks with:: + + uv run --package flashdreams --group test pytest \ + flashdreams/benchmarks/accelerated/multi_head_attention/test_multi_head_attention_benchmark.py \ + -p no:manual_marker -m manual --benchmark-only -v +""" + +from __future__ import annotations + +import pytest +import torch +from pytest_benchmark.fixture import BenchmarkFixture +from torch import Tensor + +from flashdreams.accelerated.multi_head_attention import ( + AttentionConfig, + AttentionType, + QKNormScope, + RoPEConfig, + RoPEScope, + RoPEStyle, +) +from flashdreams.accelerated.multi_head_attention.torch import TorchMultiHeadAttention +from flashdreams.accelerated.multi_head_attention.optimized import ( + QKVFusionOption, + QuantizationOption, + SDPABackend, + OptimizedImplConfig, + OptimizedHultiHeadAttention, +) + +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 implementation comparison.""" + + +_IMPLEMENTATION_CONFIGS: tuple[OptimizedImplConfig | None, ...] = ( + None, + *( + OptimizedImplConfig( + sdpa_backend=sdpa_backend, + qkv_fusion_option=qkv_fusion_option, + use_tma=use_tma, + ) + for sdpa_backend in SDPABackend + for qkv_fusion_option in QKVFusionOption + for use_tma in (False, True) + ), + *( + OptimizedImplConfig( + sdpa_backend=SDPABackend.CUDNN, + qkv_fusion_option=qkv_fusion_option, + use_tma=False, + quantization=QuantizationOption(projection=torch.float8_e4m3fn), + ) + for qkv_fusion_option in QKVFusionOption + ), + *( + OptimizedImplConfig( + sdpa_backend=sdpa_backend, + qkv_fusion_option=qkv_fusion_option, + use_tma=use_tma, + quantization=QuantizationOption( + projection=torch.float8_e4m3fn, + quantized_sdpa=True, + ), + ) + for sdpa_backend in SDPABackend + for qkv_fusion_option in QKVFusionOption + for use_tma in (False, True) + ), +) +"""Torch reference, optimized policies, FP8 projections, and quantized SDPA rows.""" + + +def _implementation_id(config: OptimizedImplConfig | None) -> str: + """Return a stable pytest identifier for an implementation config.""" + if config is None: + return "reference-torch" + backend = config.sdpa_backend.value + fusion = config.qkv_fusion_option.value.replace("_", "-") + tma = "tma" if config.use_tma else "no-tma" + projection = ( + "" + if config.quantization.projection is None + else ( + f"-projection-{str(config.quantization.projection).removeprefix('torch.')}" + ) + ) + quantized_sdpa = "-quantized-sdpa" if config.quantization.quantized_sdpa else "" + return f"optimized-{backend}-{fusion}-{tma}{projection}{quantized_sdpa}" + + +_SHARED_CONFIGS = tuple( + pytest.param( + qk_norm_scope, + rope_scope, + rope_interleaved, + bias, + id=( + f"norm-{qk_norm_scope.value}-" + f"rope-{'interleaved' if rope_interleaved else 'split'}-" + f"scope-{rope_scope.value.replace('_', '-')}-" + f"bias-{'on' if bias else 'off'}" + ), + ) + for qk_norm_scope in QKNormScope + for rope_scope in RoPEScope + for rope_interleaved in (False, True) + for bias in (False, True) +) +"""Policies shared by Torch and optimized attention, each mapped to its own benchmark group.""" + +_CROSS_ATTENTION_CONFIGS = tuple( + config + for config in _SHARED_CONFIGS + if config.values[1] is RoPEScope.BEFORE_KV_CACHE +) +"""Cross-attention policies representable with distinct query/context lengths.""" + + +class _TorchMultiHeadAttention(TorchMultiHeadAttention): + """Canonical Torch 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_bias: bool = False, + output_bias: bool = False, + qk_norm_scope: QKNormScope = QKNormScope.HEAD, + rope_scope: RoPEScope = RoPEScope.BEFORE_KV_CACHE, + rope_interleaved: bool = False, + ) -> None: + """Initialize canonical projections and normalization modules.""" + super().__init__( + attention_type=attention_type, + attention_config=AttentionConfig( + query_dim=query_dim, + context_dim=context_dim, + n_heads=n_heads, + head_dim=head_dim, + qk_norm_scope=qk_norm_scope, + rope_config=RoPEConfig( + scope=rope_scope, + style=( + RoPEStyle.INTERLEAVED if rope_interleaved else RoPEStyle.SPLIT + ), + ), + ), + ) + assert self.attention_config.context_dim is not None + self.q_proj = torch.nn.Linear( + self.attention_config.query_dim, + self.attention_config.inner_dim, + bias=qkv_bias, + ) + self.k_proj = torch.nn.Linear( + self.attention_config.context_dim, + self.attention_config.inner_dim, + bias=qkv_bias, + ) + self.v_proj = torch.nn.Linear( + self.attention_config.context_dim, + self.attention_config.inner_dim, + bias=qkv_bias, + ) + self.output_proj = torch.nn.Linear( + self.attention_config.inner_dim, + self.attention_config.query_dim, + bias=output_bias, + ) + if self.attention_config.qk_norm_scope is QKNormScope.NONE: + self.q_norm = torch.nn.Identity() + self.k_norm = torch.nn.Identity() + else: + norm_dim = ( + self.attention_config.head_dim + if self.attention_config.qk_norm_scope is QKNormScope.HEAD + else self.attention_config.inner_dim + ) + self.q_norm = torch.nn.RMSNorm( + norm_dim, eps=self.attention_config.qk_norm_eps + ) + self.k_norm = torch.nn.RMSNorm( + norm_dim, eps=self.attention_config.qk_norm_eps + ) + + +class _OptimizedHultiHeadAttention(OptimizedHultiHeadAttention): + """Canonical Optimized 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, + optimized_impl_config: OptimizedImplConfig, + qkv_bias: bool = False, + output_bias: bool = False, + qk_norm_scope: QKNormScope = QKNormScope.HEAD, + rope_scope: RoPEScope = RoPEScope.BEFORE_KV_CACHE, + rope_interleaved: bool = False, + ) -> None: + """Initialize canonical projections and normalization modules.""" + super().__init__( + attention_type=attention_type, + attention_config=AttentionConfig( + query_dim=query_dim, + context_dim=context_dim, + n_heads=n_heads, + head_dim=head_dim, + qk_norm_scope=qk_norm_scope, + rope_config=RoPEConfig( + scope=rope_scope, + style=( + RoPEStyle.INTERLEAVED if rope_interleaved else RoPEStyle.SPLIT + ), + ), + ), + optimized_impl_config=optimized_impl_config, + ) + assert self.attention_config.context_dim is not None + self.q_proj = torch.nn.Linear( + self.attention_config.query_dim, + self.attention_config.inner_dim, + bias=qkv_bias, + ) + self.k_proj = torch.nn.Linear( + self.attention_config.context_dim, + self.attention_config.inner_dim, + bias=qkv_bias, + ) + self.v_proj = torch.nn.Linear( + self.attention_config.context_dim, + self.attention_config.inner_dim, + bias=qkv_bias, + ) + self.output_proj = torch.nn.Linear( + self.attention_config.inner_dim, + self.attention_config.query_dim, + bias=output_bias, + ) + if self.attention_config.qk_norm_scope is QKNormScope.NONE: + self.q_norm = torch.nn.Identity() + self.k_norm = torch.nn.Identity() + else: + norm_dim = ( + self.attention_config.head_dim + if self.attention_config.qk_norm_scope is QKNormScope.HEAD + else self.attention_config.inner_dim + ) + self.q_norm = torch.nn.RMSNorm( + norm_dim, eps=self.attention_config.qk_norm_eps + ) + self.k_norm = torch.nn.RMSNorm( + norm_dim, eps=self.attention_config.qk_norm_eps + ) + self._initialize_derived_weights() + + +_Attention = _TorchMultiHeadAttention | _OptimizedHultiHeadAttention + +_BATCH_SIZE = 1 +_DTYPE = torch.bfloat16 +_SEED = 42 +_SINK_SIZE = 0 + + +_QUERY_DIM = 2048 +"""Input and output feature width shared by both implementations.""" + +_N_HEADS = 16 +"""Number of attention heads shared by both implementations.""" + +_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( + optimized_impl_config: OptimizedImplConfig | None, + *, + attention_type: AttentionType, + qk_norm_scope: QKNormScope, + rope_scope: RoPEScope, + rope_interleaved: bool, + bias: bool, +) -> _Attention: + """Build one weight-matched Torch reference or Optimized implementation. + + Args: + optimized_impl_config: Optimized policy; ``None`` uses the Torch reference. + attention_type: Whether to configure streaming self-attention or static + cross-attention. + qk_norm_scope: Shared Q/K normalization policy. + rope_scope: Whether keys are rotated before or after cache storage. + rope_interleaved: Shared rotary-pair layout. + bias: Whether every Q/K/V and output projection uses a bias. + + Returns: + Configured attention module with deterministic random weights. + """ + reference = _TorchMultiHeadAttention( + query_dim=_QUERY_DIM, + context_dim=_QUERY_DIM, + n_heads=_N_HEADS, + head_dim=_HEAD_DIM, + attention_type=attention_type, + qkv_bias=bias, + output_bias=bias, + qk_norm_scope=qk_norm_scope, + rope_scope=rope_scope, + rope_interleaved=rope_interleaved, + ) + if optimized_impl_config is None: + return reference + + optimized_attention = _OptimizedHultiHeadAttention( + query_dim=_QUERY_DIM, + context_dim=_QUERY_DIM, + n_heads=_N_HEADS, + head_dim=_HEAD_DIM, + attention_type=attention_type, + optimized_impl_config=optimized_impl_config, + qkv_bias=bias, + output_bias=bias, + qk_norm_scope=qk_norm_scope, + rope_scope=rope_scope, + rope_interleaved=rope_interleaved, + ) + optimized_attention.load_state_dict(reference.state_dict(), strict=True) + return optimized_attention + + +@torch.inference_mode() +def _benchmark_multi_head_attention( + benchmark: BenchmarkFixture, + optimized_impl_config: OptimizedImplConfig | None, + *, + attention_type: AttentionType, + qk_norm_scope: QKNormScope, + rope_scope: RoPEScope, + 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. + optimized_impl_config: Optimized policy; ``None`` uses the Torch reference. + attention_type: Self- or cross-attention benchmark family. + qk_norm_scope: Shared Q/K normalization policy. + rope_scope: Whether keys are rotated before or after cache storage. + 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") + is_optimized = optimized_impl_config is not None + if is_optimized and torch.cuda.get_device_capability(device)[0] < 9: + pytest.skip("Optimized attention requires compute capability 9.0 or newer.") + + torch.manual_seed(_SEED) + attention = _make_attention( + optimized_impl_config, + attention_type=attention_type, + qk_norm_scope=qk_norm_scope, + rope_scope=rope_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_freq_count = ( + 1 if rope_scope is RoPEScope.AFTER_KV_CACHE else _WINDOW_CHUNKS + 1 + ) + rope_freq_length = ( + _WINDOW_SIZE if rope_scope is RoPEScope.AFTER_KV_CACHE else _CHUNK_SIZE + ) + half_rope_freqs = [ + torch.randn( + rope_freq_length, + 1, + 1, + _HEAD_DIM // 2, + generator=generator, + device=device, + dtype=torch.float32, + ) + for _ in range(rope_freq_count) + ] + rope_freqs = [ + ( + freqs.repeat_interleave(2, dim=-1) + if rope_interleaved + else torch.cat((freqs, freqs), dim=-1) + ) + for freqs in half_rope_freqs + ] + if rope_scope is RoPEScope.AFTER_KV_CACHE: + rope_freqs *= _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, + "scope", + rope_scope.value.replace("_", "-"), + "bias", + bias_label, + ) + ) + benchmark.extra_info.update( + { + "gpu": torch.cuda.get_device_name(device), + "torch": str(torch.__version__), + "cuda": torch.version.cuda, + "cudnn": torch.backends.cudnn.version(), + "input_dtype": str(_DTYPE), + "projection_dtype": ( + None + if optimized_impl_config is None + or optimized_impl_config.quantization.projection is None + else str(optimized_impl_config.quantization.projection) + ), + "quantized_sdpa": ( + False + if optimized_impl_config is None + else optimized_impl_config.quantization.quantized_sdpa + ), + "sdpa_backend": ( + None + if optimized_impl_config is None + else optimized_impl_config.sdpa_backend.value + ), + "qkv_fusion_option": ( + None + if optimized_impl_config is None + else optimized_impl_config.qkv_fusion_option.value + ), + "use_tma": ( + None if optimized_impl_config is None else optimized_impl_config.use_tma + ), + "seed": _SEED, + "batch_size": _BATCH_SIZE, + "chunk_size": _CHUNK_SIZE, + "window_size": _WINDOW_SIZE, + "warmup_rounds": _WARMUP_ROUNDS, + "benchmark_rounds": _BENCHMARK_ROUNDS, + } + ) + + 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( + "optimized_impl_config", + _IMPLEMENTATION_CONFIGS, + ids=_implementation_id, +) +@pytest.mark.parametrize( + "qk_norm_scope,rope_scope,rope_interleaved,bias", + _SHARED_CONFIGS, +) +def test_self_attention_benchmark( + benchmark: BenchmarkFixture, + optimized_impl_config: OptimizedImplConfig | None, + qk_norm_scope: QKNormScope, + rope_scope: RoPEScope, + rope_interleaved: bool, + bias: bool, +) -> None: + """Compare streaming self-attention within one shared-policy group. + + Args: + benchmark: Pytest benchmark fixture used to record synchronized timings. + optimized_impl_config: Optimized policy; ``None`` uses the Torch reference. + qk_norm_scope: Shared Q/K normalization policy defining the group. + rope_scope: Shared cache-relative rotation 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, + optimized_impl_config, + attention_type=AttentionType.SELF_ATTENTION, + qk_norm_scope=qk_norm_scope, + rope_scope=rope_scope, + rope_interleaved=rope_interleaved, + bias=bias, + ) + + +@pytest.mark.parametrize( + "optimized_impl_config", + _IMPLEMENTATION_CONFIGS, + ids=_implementation_id, +) +@pytest.mark.parametrize( + "qk_norm_scope,rope_scope,rope_interleaved,bias", + _CROSS_ATTENTION_CONFIGS, +) +def test_cross_attention_benchmark( + benchmark: BenchmarkFixture, + optimized_impl_config: OptimizedImplConfig | None, + qk_norm_scope: QKNormScope, + rope_scope: RoPEScope, + rope_interleaved: bool, + bias: bool, +) -> None: + """Compare end-to-end cross-attention within one shared-policy group. + + Args: + benchmark: Pytest benchmark fixture used to record synchronized timings. + optimized_impl_config: Optimized policy; ``None`` uses the Torch reference. + qk_norm_scope: Shared Q/K normalization policy defining the group. + rope_scope: Shared cache-relative rotation 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, + optimized_impl_config, + attention_type=AttentionType.CROSS_ATTENTION, + qk_norm_scope=qk_norm_scope, + rope_scope=rope_scope, + rope_interleaved=rope_interleaved, + bias=bias, + ) diff --git a/flashdreams/benchmarks/accelerated/quantization/test_quantized_gemm_benchmark.py b/flashdreams/benchmarks/accelerated/quantization/test_quantized_gemm_benchmark.py new file mode 100644 index 000000000..cb1fc5df5 --- /dev/null +++ b/flashdreams/benchmarks/accelerated/quantization/test_quantized_gemm_benchmark.py @@ -0,0 +1,263 @@ +# 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. + +"""Benchmarks for full-precision and quantized Torch matrix multiplication. + +Quantized ``end-to-end`` rows include operand quantization and output conversion +in the timed region. ``gemm-only`` rows prepare operands before timing and omit +INT8 dequantization. FP8 scale application remains fused into ``torch._scaled_mm``. +CUDA does not support E5M2 by E5M2 GEMM, so E5M2 rows use E5M2 for the left +operand and E4M3 for the right operand. + +Run the manual GPU benchmarks with:: + + uv run --package flashdreams --group test pytest \ + flashdreams/benchmarks/accelerated/quantization/test_quantized_gemm_benchmark.py \ + -p no:manual_marker -m manual --benchmark-only -v +""" + +from __future__ import annotations + +from collections.abc import Callable + +import pytest +import torch +from pytest_benchmark.fixture import BenchmarkFixture +from torch import Tensor + +from flashdreams.accelerated.quantization.quantizer import ( + DTYPE_MAX, + Granularity, + dequantize, + quantize, +) + +pytestmark = [ + pytest.mark.manual, + pytest.mark.skipif( + not torch.cuda.is_available(), + reason="Quantized GEMM benchmarks require CUDA.", + ), +] + +_M = 4096 +_K = 4096 +_N = 4096 +_SEED = 42 +_WARMUP_ROUNDS = 5 +"""Warmup calls used to absorb kernel initialization and autotuning.""" + +_BENCHMARK_ROUNDS = 50 +"""Measured calls used for each GEMM comparison.""" + +_ORIGINAL_DTYPES = ( + pytest.param(torch.float16, "fp16", id="fp16"), + pytest.param(torch.bfloat16, "bf16", id="bf16"), + pytest.param(torch.float32, "fp32", id="fp32"), +) +"""Original matrix formats, each mapped to a separate benchmark group.""" + +_GEMM_CASES = ( + pytest.param(None, None, id="full-precision"), + *( + pytest.param( + quantized_dtype, + granularity, + id=( + f"{str(quantized_dtype).removeprefix('torch.')}" + f"{'-x-float8_e4m3fn' if quantized_dtype is torch.float8_e5m2 else ''}" + f"-{granularity.value}" + ), + ) + for quantized_dtype in DTYPE_MAX + for granularity in Granularity + ), +) +"""Full-precision baseline and every supported quantized GEMM configuration.""" + + +def _quantized_gemm( + left: Tensor, + right: Tensor, + quantized_dtype: torch.dtype, + granularity: Granularity, + output_dtype: torch.dtype, + end_to_end: bool, +) -> tuple[Callable[[], Tensor], torch.dtype]: + """Return the timed operation and the data type it produces.""" + if quantized_dtype is torch.int8: + + def quantize_operands() -> tuple[Tensor, Tensor, Tensor, Tensor]: + left_quantized, left_scale = quantize( + left, quantized_dtype, granularity, axis=1 + ) + right_quantized, right_scale = quantize( + right, quantized_dtype, granularity, axis=0 + ) + return left_quantized, right_quantized, left_scale, right_scale + + prepared_operands = None if end_to_end else quantize_operands() + + def gemm() -> Tensor: + operands = quantize_operands() if end_to_end else prepared_operands + assert operands is not None + left_quantized, right_quantized, left_scale, right_scale = operands + output = torch._int_mm(left_quantized, right_quantized) + if end_to_end: + return dequantize( + output, + left_scale, + right_scale, + dtype=output_dtype, + ) + return output + + return gemm, output_dtype if end_to_end else torch.int32 + + right_dtype = ( + torch.float8_e4m3fn if quantized_dtype is torch.float8_e5m2 else quantized_dtype + ) + scaled_output_dtype = ( + torch.bfloat16 + if granularity is Granularity.SLICE and output_dtype is torch.float32 + else output_dtype + ) + + def quantize_operands() -> tuple[Tensor, Tensor, Tensor, Tensor]: + left_quantized, left_scale = quantize( + left, quantized_dtype, granularity, axis=1 + ) + right_quantized_transposed, right_scale_transposed = quantize( + right.t().contiguous(), right_dtype, granularity, axis=1 + ) + return ( + left_quantized, + right_quantized_transposed.t(), + left_scale, + right_scale_transposed.t(), + ) + + prepared_operands = None if end_to_end else quantize_operands() + + def gemm() -> Tensor: + operands = quantize_operands() if end_to_end else prepared_operands + assert operands is not None + left_quantized, right_quantized, left_scale, right_scale = operands + output = torch._scaled_mm( + left_quantized, + right_quantized, + left_scale, + right_scale, + out_dtype=scaled_output_dtype, + ) + return output.to(output_dtype) if end_to_end else output + + return gemm, output_dtype if end_to_end else scaled_output_dtype + + +def _benchmark_quantized_gemm( + benchmark: BenchmarkFixture, + original_dtype: torch.dtype, + original_format: str, + quantized_dtype: torch.dtype | None, + granularity: Granularity | None, + *, + end_to_end: bool, +) -> None: + """Benchmark one quantized GEMM configuration.""" + generator = torch.Generator(device="cuda").manual_seed(_SEED) + left = torch.randn( + (_M, _K), device="cuda", dtype=original_dtype, generator=generator + ) + right = torch.randn( + (_K, _N), device="cuda", dtype=original_dtype, generator=generator + ) + + if quantized_dtype is None: + output_dtype = original_dtype + + def operation() -> Tensor: + return left @ right + + else: + assert granularity is not None + operation, output_dtype = _quantized_gemm( + left, + right, + quantized_dtype, + granularity, + original_dtype, + end_to_end, + ) + + timing_scope = "end-to-end" if end_to_end else "gemm-only" + benchmark.group = f"quantized-gemm-{original_format}-{timing_scope}" + + def synchronized_gemm() -> Tensor: + output = operation() + torch.cuda.synchronize() + return output + + torch.cuda.synchronize() + output = benchmark.pedantic( + synchronized_gemm, + iterations=1, + rounds=_BENCHMARK_ROUNDS, + warmup_rounds=_WARMUP_ROUNDS, + ) + + assert output.shape == (_M, _N) + assert output.dtype is output_dtype + assert torch.isfinite(output).all() + + +@pytest.mark.parametrize("quantized_dtype,granularity", _GEMM_CASES) +@pytest.mark.parametrize("original_dtype,original_format", _ORIGINAL_DTYPES) +def test_quantized_gemm_end_to_end_benchmark( + benchmark: BenchmarkFixture, + original_dtype: torch.dtype, + original_format: str, + quantized_dtype: torch.dtype | None, + granularity: Granularity | None, +) -> None: + """Benchmark quantization, GEMM, and output conversion together.""" + _benchmark_quantized_gemm( + benchmark, + original_dtype, + original_format, + quantized_dtype, + granularity, + end_to_end=True, + ) + + +@pytest.mark.parametrize("quantized_dtype,granularity", _GEMM_CASES) +@pytest.mark.parametrize("original_dtype,original_format", _ORIGINAL_DTYPES) +def test_quantized_gemm_only_benchmark( + benchmark: BenchmarkFixture, + original_dtype: torch.dtype, + original_format: str, + quantized_dtype: torch.dtype | None, + granularity: Granularity | None, +) -> None: + """Benchmark GEMM with operands quantized before timing.""" + _benchmark_quantized_gemm( + benchmark, + original_dtype, + original_format, + quantized_dtype, + granularity, + end_to_end=False, + ) diff --git a/flashdreams/benchmarks/accelerated/quantization/test_quantized_linear_benchmark.py b/flashdreams/benchmarks/accelerated/quantization/test_quantized_linear_benchmark.py new file mode 100644 index 000000000..8b8b2db12 --- /dev/null +++ b/flashdreams/benchmarks/accelerated/quantization/test_quantized_linear_benchmark.py @@ -0,0 +1,224 @@ +# 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. + +"""Benchmarks for regular and quantized nonpersistent linear inference. + +``full-precision-x`` rows include activation quantization in the timed region, +while ``prequantized-x`` rows prepare activations and scales before timing. +Weight quantization always happens during module construction. + +Run the manual GPU benchmarks with:: + + uv run --package flashdreams --group test pytest \ + flashdreams/benchmarks/accelerated/quantization/test_quantized_linear_benchmark.py \ + -p no:manual_marker -m manual --benchmark-only -v +""" + +from collections.abc import Callable + +import pytest +import torch +import torch.nn.functional as F +from pytest_benchmark.fixture import BenchmarkFixture +from torch import Tensor, nn + +from flashdreams.accelerated.quantization.linear import ( + QuantizedNonPersistentLinear, + WeightGranularity, +) +from flashdreams.accelerated.quantization.quantizer import ( + DTYPE_MAX, + Granularity, + quantize, +) + +pytestmark = [ + pytest.mark.manual, + pytest.mark.skipif( + not torch.cuda.is_available(), + reason="Quantized linear benchmarks require CUDA.", + ), +] + +_BATCH_SIZE = 4096 +_IN_FEATURES = 4096 +_OUT_FEATURES = 4096 +_SEED = 42 +_WARMUP_ROUNDS = 5 +"""Warmup calls used to absorb kernel initialization and autotuning.""" + +_BENCHMARK_ROUNDS = 50 +"""Measured calls used for each linear comparison.""" + +_ORIGINAL_DTYPES = ( + pytest.param(torch.float16, "fp16", id="fp16"), + pytest.param(torch.bfloat16, "bf16", id="bf16"), + pytest.param(torch.float32, "fp32", id="fp32"), +) +"""Original activation and output formats, each in a separate benchmark group.""" + +_LINEAR_CASES = ( + pytest.param(None, None, None, None, id="nn-linear"), + *( + pytest.param( + quantized_dtype, + weight_granularity, + input_granularity, + prequantized, + id=( + f"{str(quantized_dtype).removeprefix('torch.')}" + f"{'-x-float8_e4m3fn' if quantized_dtype is torch.float8_e5m2 else ''}" + f"-weight-{weight_granularity.value}" + f"-input-{input_granularity.value}" + f"-{'prequantized-x' if prequantized else 'full-precision-x'}" + ), + ) + for quantized_dtype in DTYPE_MAX + for weight_granularity in WeightGranularity + for input_granularity in Granularity + for prequantized in (False, True) + ), +) +"""Regular linear baseline and every quantized linear inference configuration.""" + + +@pytest.mark.parametrize( + "quantized_dtype,weight_granularity,input_granularity,prequantized", + _LINEAR_CASES, +) +@pytest.mark.parametrize("original_dtype,original_format", _ORIGINAL_DTYPES) +@torch.inference_mode() +def test_quantized_linear_benchmark( + benchmark: BenchmarkFixture, + original_dtype: torch.dtype, + original_format: str, + quantized_dtype: torch.dtype | None, + weight_granularity: WeightGranularity | None, + input_granularity: Granularity | None, + prequantized: bool | None, +) -> None: + """Benchmark regular and quantized linear inference.""" + generator = torch.Generator(device="cuda").manual_seed(_SEED) + inputs = torch.randn( + (_BATCH_SIZE, _IN_FEATURES), + device="cuda", + dtype=original_dtype, + generator=generator, + ) + weight = torch.randn( + (_OUT_FEATURES, _IN_FEATURES), + device="cuda", + dtype=original_dtype, + generator=generator, + ) + + operation: Callable[[], Tensor] + if quantized_dtype is None: + assert ( + weight_granularity is None + and input_granularity is None + and prequantized is None + ) + linear = nn.Linear( + _IN_FEATURES, + _OUT_FEATURES, + bias=False, + device="cuda", + dtype=original_dtype, + ).requires_grad_(False) + linear.weight.copy_(weight) + + def operation() -> Tensor: + return linear(inputs) + + else: + assert weight_granularity is not None + assert input_granularity is not None + assert prequantized is not None + quantized_linear = QuantizedNonPersistentLinear( + weight, + None, + weight_granularity, + quantized_dtype, + ) + if prequantized: + quantized_inputs, input_scale = quantize( + inputs, + quantized_dtype, + input_granularity, + axis=-1, + ) + + def operation() -> Tensor: + return quantized_linear( + quantized_inputs, + input_scale, + out_dtype=original_dtype, + ) + + else: + + def operation() -> Tensor: + return quantized_linear( + inputs, + input_granularity, + out_dtype=original_dtype, + ) + + benchmark.group = f"quantized-linear-{original_format}" + benchmark.extra_info.update( + { + "gpu": torch.cuda.get_device_name(), + "torch": torch.__version__, + "cuda": torch.version.cuda, + "batch_size": _BATCH_SIZE, + "in_features": _IN_FEATURES, + "out_features": _OUT_FEATURES, + "warmup_rounds": _WARMUP_ROUNDS, + "benchmark_rounds": _BENCHMARK_ROUNDS, + } + ) + + def synchronized_linear() -> Tensor: + output = operation() + torch.cuda.synchronize() + return output + + torch.cuda.synchronize() + output = benchmark.pedantic( + synchronized_linear, + iterations=1, + rounds=_BENCHMARK_ROUNDS, + warmup_rounds=_WARMUP_ROUNDS, + ) + + assert output.shape == (_BATCH_SIZE, _OUT_FEATURES) + assert output.dtype is original_dtype + assert torch.isfinite(output).all() + + reference = F.linear(inputs, weight) + relative_error = ( + output.float() - reference.float() + ).norm() / reference.float().norm() + tolerance = ( + 0.0 + if quantized_dtype is None + else ( + torch.finfo(quantized_dtype).eps + if quantized_dtype.is_floating_point + else 4 / DTYPE_MAX[quantized_dtype] + ) + ) + assert relative_error.item() <= tolerance diff --git a/flashdreams/benchmarks/accelerated/quantization/test_quantizer_benchmark.py b/flashdreams/benchmarks/accelerated/quantization/test_quantizer_benchmark.py new file mode 100644 index 000000000..d828b3fd3 --- /dev/null +++ b/flashdreams/benchmarks/accelerated/quantization/test_quantizer_benchmark.py @@ -0,0 +1,189 @@ +# 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. + +"""Benchmarks comparing Torch and Triton tensor quantization. + +Run the manual GPU benchmarks with:: + + uv run --package flashdreams --group test pytest \ + flashdreams/benchmarks/accelerated/quantization/test_quantizer_benchmark.py \ + -p no:manual_marker -m manual --benchmark-only -v +""" + +from __future__ import annotations + +import pytest +import torch +import triton +from pytest_benchmark.fixture import BenchmarkFixture +from torch import Tensor + +from flashdreams.accelerated.quantization.quantizer import ( + DTYPE_MAX, + Granularity, + dequantize, + quantize, +) + +pytestmark = [ + pytest.mark.manual, + pytest.mark.skipif( + not torch.cuda.is_available(), + reason="Quantizer benchmarks require CUDA.", + ), +] + +_SHAPE = (4096, 4096) +"""Matrix shape representative of large projection activations and weights.""" + +_SEED = 42 + +_WARMUP_ROUNDS = 5 +"""Warmup calls used to absorb Triton compilation and CUDA initialization.""" + +_BENCHMARK_ROUNDS = 50 +"""Measured calls used for each Torch and Triton comparison.""" + +_IMPLEMENTATIONS = ( + pytest.param(False, "torch", id="torch"), + pytest.param(True, "triton", id="triton"), +) + + +def _record_environment( + benchmark: BenchmarkFixture, + format: torch.dtype, + granularity: Granularity, + implementation: str, +) -> None: + """Attach the reproducibility metadata shared by both operations.""" + benchmark.extra_info.update( + { + "gpu": torch.cuda.get_device_name(), + "torch": torch.__version__, + "triton": triton.__version__, + "cuda": torch.version.cuda, + "shape": _SHAPE, + "source_dtype": str(torch.float16), + "quantized_dtype": str(format), + "granularity": granularity.value, + "implementation": implementation, + "warmup_rounds": _WARMUP_ROUNDS, + "benchmark_rounds": _BENCHMARK_ROUNDS, + } + ) + + +@pytest.mark.parametrize("format", DTYPE_MAX) +@pytest.mark.parametrize("granularity", Granularity) +@pytest.mark.parametrize("use_triton,implementation", _IMPLEMENTATIONS) +def test_quantize_benchmark( + benchmark: BenchmarkFixture, + format: torch.dtype, + granularity: Granularity, + use_triton: bool, + implementation: str, +) -> None: + """Benchmark quantization through one Torch or Triton implementation.""" + generator = torch.Generator(device="cuda").manual_seed(_SEED) + original = torch.randn( + _SHAPE, + device="cuda", + dtype=torch.float16, + generator=generator, + ) + format_name = str(format).removeprefix("torch.") + benchmark.group = f"quantize-{format_name}-{granularity.value}" + _record_environment(benchmark, format, granularity, implementation) + + def synchronized_quantize() -> tuple[Tensor, Tensor]: + output = quantize( + original, + format, + granularity, + axis=-1, + use_triton=use_triton, + ) + torch.cuda.synchronize() + return output + + torch.cuda.synchronize() + quantized, scale = benchmark.pedantic( + synchronized_quantize, + iterations=1, + rounds=_BENCHMARK_ROUNDS, + warmup_rounds=_WARMUP_ROUNDS, + ) + + expected_scale_shape = ( + (1, 1) if granularity is Granularity.TENSOR else (_SHAPE[0], 1) + ) + assert quantized.shape == _SHAPE + assert quantized.dtype is format + assert scale.shape == expected_scale_shape + assert torch.isfinite(quantized.float()).all() + assert torch.isfinite(scale).all() + + +@pytest.mark.parametrize("format", DTYPE_MAX) +@pytest.mark.parametrize("granularity", Granularity) +@pytest.mark.parametrize("use_triton,implementation", _IMPLEMENTATIONS) +def test_dequantize_benchmark( + benchmark: BenchmarkFixture, + format: torch.dtype, + granularity: Granularity, + use_triton: bool, + implementation: str, +) -> None: + """Benchmark dequantization through one Torch or Triton implementation.""" + generator = torch.Generator(device="cuda").manual_seed(_SEED) + original = torch.randn( + _SHAPE, + device="cuda", + dtype=torch.float16, + generator=generator, + ) + quantized, scale = quantize( + original, + format, + granularity, + axis=-1, + use_triton=False, + ) + format_name = str(format).removeprefix("torch.") + benchmark.group = f"dequantize-{format_name}-{granularity.value}" + _record_environment(benchmark, format, granularity, implementation) + + def synchronized_dequantize() -> Tensor: + output = dequantize( + quantized, + scale, + dtype=torch.float16, + use_triton=use_triton, + ) + torch.cuda.synchronize() + return output + + torch.cuda.synchronize() + dequantized = benchmark.pedantic( + synchronized_dequantize, + iterations=1, + rounds=_BENCHMARK_ROUNDS, + warmup_rounds=_WARMUP_ROUNDS, + ) + + assert dequantized.shape == _SHAPE + assert dequantized.dtype is torch.float16 + assert torch.isfinite(dequantized).all() diff --git a/flashdreams/flashdreams/accelerated/common/__init__.py b/flashdreams/flashdreams/accelerated/common/__init__.py new file mode 100644 index 000000000..ddacfa8a8 --- /dev/null +++ b/flashdreams/flashdreams/accelerated/common/__init__.py @@ -0,0 +1,20 @@ +# 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 modules for accelerated inference implementations.""" + +from flashdreams.accelerated.common.non_persistent_linear import NonPersistentLinear + +__all__ = ["NonPersistentLinear"] diff --git a/flashdreams/flashdreams/accelerated/common/non_persistent_linear.py b/flashdreams/flashdreams/accelerated/common/non_persistent_linear.py new file mode 100644 index 000000000..11bdf4a74 --- /dev/null +++ b/flashdreams/flashdreams/accelerated/common/non_persistent_linear.py @@ -0,0 +1,35 @@ +# 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. + +"""Nonpersistent linear transformation for accelerated inference.""" + +from torch import Tensor, nn + + +class NonPersistentLinear(nn.Linear): + """Linear transformation backed by nonpersistent weight and bias buffers.""" + + def __init__(self, weight: Tensor, bias: Tensor | None) -> None: + """Initialize the transformation from existing tensors. + + Args: + weight: Projection weight shaped ``[out_features, in_features]``. + bias: Optional projection bias shaped ``[out_features]``. + """ + nn.Module.__init__(self) + self.in_features = weight.shape[1] + self.out_features = weight.shape[0] + self.register_buffer("weight", weight, persistent=False) + self.register_buffer("bias", bias, persistent=False) diff --git a/flashdreams/flashdreams/accelerated/multi_head_attention/__init__.py b/flashdreams/flashdreams/accelerated/multi_head_attention/__init__.py new file mode 100644 index 000000000..ca8766268 --- /dev/null +++ b/flashdreams/flashdreams/accelerated/multi_head_attention/__init__.py @@ -0,0 +1,306 @@ +# 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 dataclasses import dataclass +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 RoPEStyle(str, Enum): + """Feature pairing convention for rotary position embeddings.""" + + INTERLEAVED = "interleaved" + """Rotate adjacent feature pairs.""" + + SPLIT = "split" + """Rotate corresponding features from the two half-width blocks.""" + + +class RoPEScope(str, Enum): + """Point at which rotary position embeddings are applied to keys.""" + + BEFORE_KV_CACHE = "before_kv_cache" + """Rotate keys before writing them to the K/V cache.""" + + AFTER_KV_CACHE = "after_kv_cache" + """Rotate cached keys immediately before attention.""" + + +@dataclass(frozen=True, slots=True) +class RoPEConfig: + """Rotary position embedding policy.""" + + style: RoPEStyle = RoPEStyle.SPLIT + """Feature pairing convention used by each rotary embedding.""" + + scope: RoPEScope = RoPEScope.BEFORE_KV_CACHE + """Whether keys are rotated before or after K/V cache storage.""" + + def __post_init__(self) -> None: + """Validate rotary policy enum values.""" + if not isinstance(self.style, RoPEStyle): + raise TypeError(f"style must be a RoPEStyle; got {self.style!r}") + if not isinstance(self.scope, RoPEScope): + raise TypeError(f"scope must be a RoPEScope; got {self.scope!r}") + + +@dataclass(frozen=True, slots=True) +class AttentionConfig: + """Geometry, normalization, and rotary policy shared by attention backends.""" + + query_dim: int + """Input and output token width.""" + + n_heads: int = 8 + """Number of query, key, and value heads.""" + + head_dim: int = 64 + """Feature width of each attention head.""" + + context_dim: int | None = None + """Context token width; ``None`` uses ``query_dim``.""" + + qk_norm_scope: QKNormScope = QKNormScope.HEAD + """Feature scope used by query and key normalization.""" + + qk_norm_eps: float = 1e-6 + """Positive finite epsilon used by query and key normalization.""" + + rope_config: RoPEConfig | None = None + """Rotary policy; ``None`` disables rotary embeddings.""" + + @property + def inner_dim(self) -> int: + """Return the concatenated width of all attention heads.""" + return self.n_heads * self.head_dim + + def __post_init__(self) -> None: + """Validate and normalize attention configuration values.""" + context_dim = self.query_dim if self.context_dim is None else self.context_dim + if self.query_dim <= 0: + raise ValueError(f"query_dim must be positive; got {self.query_dim}") + if context_dim <= 0: + raise ValueError(f"context_dim must be positive; got {context_dim}") + if self.n_heads <= 0: + raise ValueError(f"n_heads must be positive; got {self.n_heads}") + if self.head_dim <= 0: + raise ValueError(f"head_dim must be positive; got {self.head_dim}") + if not isinstance(self.qk_norm_scope, QKNormScope): + raise TypeError( + f"qk_norm_scope must be a QKNormScope; got {self.qk_norm_scope!r}" + ) + if not math.isfinite(self.qk_norm_eps) or self.qk_norm_eps <= 0: + raise ValueError( + f"qk_norm_eps must be finite and positive; got {self.qk_norm_eps}" + ) + if self.rope_config is not None and not isinstance( + self.rope_config, RoPEConfig + ): + raise TypeError( + f"rope_config must be a RoPEConfig or None; got {self.rope_config!r}" + ) + object.__setattr__(self, "context_dim", context_dim) + + +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. + """ + + attention_type: AttentionType + """Whether forward performs self-attention or static cross-attention.""" + + attention_config: AttentionConfig + """Geometry, normalization, and rotary policy used by this module.""" + + def __init__( + self, + attention_type: AttentionType, + attention_config: AttentionConfig, + ) -> None: + """Initialize shared attention geometry and implementation policies. + + Args: + attention_type: Select self-attention or static cross-attention. + attention_config: Shared geometry, normalization, and rotary policy. + + Raises: + TypeError: An argument has the wrong configuration type. + ValueError: Self-attention has different query and context dimensions. + """ + super().__init__() + + if not isinstance(attention_type, AttentionType): + raise TypeError( + f"attention_type must be an AttentionType; got {attention_type!r}" + ) + if not isinstance(attention_config, AttentionConfig): + raise TypeError( + f"attention_config must be an AttentionConfig; got {attention_config!r}" + ) + context_dim = attention_config.context_dim + assert context_dim is not None + if ( + attention_type is AttentionType.SELF_ATTENTION + and attention_config.query_dim != context_dim + ): + raise ValueError( + "self-attention requires query_dim to equal context_dim; " + f"got {attention_config.query_dim} and {context_dim}" + ) + + self.attention_type = attention_type + self.attention_config = attention_config + + @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. Applied only by before-cache RoPE; after-cache RoPE + stores unrotated keys and consumes cache-relative data in + :meth:`forward`. Ignored when ``rope_config`` is ``None``. + + 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 positional data. Before-cache RoPE expects the + current ``L`` query/key positions. After-cache RoPE expects + cache-relative positions for all visible keys and selects the + current query positions from the cache write interval. Ignored + when ``rope_config`` is ``None``. + + Returns: + Attention result with shape ``[..., L, query_dim]``. + """ + + +__all__ = [ + "AttentionConfig", + "AttentionType", + "MultiHeadAttention", + "QKNormScope", + "RoPEConfig", + "RoPEScope", + "RoPEStyle", +] diff --git a/flashdreams/flashdreams/accelerated/multi_head_attention/cudnn/__init__.py b/flashdreams/flashdreams/accelerated/multi_head_attention/cudnn/__init__.py new file mode 100644 index 000000000..8dbb4990a --- /dev/null +++ b/flashdreams/flashdreams/accelerated/multi_head_attention/cudnn/__init__.py @@ -0,0 +1,25 @@ +# 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. + +"""Torch and native cuDNN scaled-dot-product attention.""" + +from flashdreams.accelerated.multi_head_attention.cudnn.native_fp8 import ( + native_cudnn_fp8_sdpa, +) +from flashdreams.accelerated.multi_head_attention.cudnn.torch_sdpa import ( + torch_cudnn_sdpa, +) + +__all__ = ["native_cudnn_fp8_sdpa", "torch_cudnn_sdpa"] diff --git a/flashdreams/flashdreams/accelerated/multi_head_attention/cudnn/native_fp8.py b/flashdreams/flashdreams/accelerated/multi_head_attention/cudnn/native_fp8.py new file mode 100644 index 000000000..78fedd466 --- /dev/null +++ b/flashdreams/flashdreams/accelerated/multi_head_attention/cudnn/native_fp8.py @@ -0,0 +1,158 @@ +# 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. + +"""Native cuDNN Frontend FP8 scaled-dot-product attention.""" + +from __future__ import annotations + +import importlib +import math +from collections.abc import Callable +from functools import lru_cache + +import torch +from torch import Tensor + + +@lru_cache(maxsize=256) +def _build_cudnn_fp8_sdpa( + device: torch.device, + stream: int, + query_shape: tuple[int, ...], + query_stride: tuple[int, ...], + key_shape: tuple[int, ...], + key_stride: tuple[int, ...], + value_shape: tuple[int, ...], + value_stride: tuple[int, ...], +) -> Callable[[Tensor, Tensor, Tensor], Tensor]: + """Build one shape-, layout-, device-, and stream-specialized FP8 graph.""" + try: + cudnn = importlib.import_module("cudnn") + except ModuleNotFoundError as error: + raise RuntimeError( + "quantized cuDNN SDPA requires nvidia-cudnn-frontend" + ) from error + + graph = cudnn.pygraph( + io_data_type=cudnn.data_type.FP8_E4M3, + intermediate_data_type=cudnn.data_type.FLOAT, + compute_data_type=cudnn.data_type.FLOAT, + ) + + def tensor(name: str, shape: tuple[int, ...], stride: tuple[int, ...]): + return graph.tensor( + name=name, + dim=list(shape), + stride=list(stride), + data_type=cudnn.data_type.FP8_E4M3, + ) + + query_desc = tensor("query", query_shape, query_stride) + key_desc = tensor("key", key_shape, key_stride) + value_desc = tensor("value", value_shape, value_stride) + scale_descriptors = tuple( + graph.tensor( + name=name, + dim=[1, 1, 1, 1], + stride=[1, 1, 1, 1], + data_type=cudnn.data_type.FLOAT, + ) + for name in ( + "descale_q", + "descale_k", + "descale_v", + "descale_s", + "scale_s", + "scale_o", + ) + ) + output_desc, _, amax_s_desc, amax_o_desc = graph.sdpa_fp8( + q=query_desc, + k=key_desc, + v=value_desc, + descale_q=scale_descriptors[0], + descale_k=scale_descriptors[1], + descale_v=scale_descriptors[2], + descale_s=scale_descriptors[3], + scale_s=scale_descriptors[4], + scale_o=scale_descriptors[5], + is_inference=True, + attn_scale=1.0 / math.sqrt(query_shape[-1]), + name="sdpa", + ) + + output = torch.empty_strided( + query_shape, query_stride, dtype=torch.float8_e4m3fn, device=device + ) + amax_s = torch.empty((1, 1, 1, 1), dtype=torch.float32, device=device) + amax_o = torch.empty_like(amax_s) + output_desc.set_output(True).set_dim(list(output.shape)).set_stride( + list(output.stride()) + ) + amax_s_desc.set_output(False).set_dim(list(amax_s.shape)).set_stride( + list(amax_s.stride()) + ) + amax_o_desc.set_output(False).set_dim(list(amax_o.shape)).set_stride( + list(amax_o.stride()) + ) + graph.build([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) + scale = torch.ones_like(amax_s) + workspace = torch.empty( + graph.get_workspace_size(), dtype=torch.uint8, device=device + ) + + def execute(query: Tensor, key: Tensor, value: Tensor) -> Tensor: + graph.execute( + { + query_desc: query, + key_desc: key, + value_desc: value, + **{descriptor: scale for descriptor in scale_descriptors}, + output_desc: output, + amax_s_desc: amax_s, + amax_o_desc: amax_o, + }, + workspace, + ) + return output + + return execute + + +def native_cudnn_fp8_sdpa(query: Tensor, key: Tensor, value: Tensor) -> Tensor: + """Apply unscaled e4m3 attention with a cached cuDNN Frontend graph. + + Args: + query: FP8 queries in ``[B, H, L, D]`` layout. + key: FP8 keys in ``[B, H, S, D]`` layout. + value: FP8 values in ``[B, H, S, D]`` layout. + + Returns: + FP8 attention output shaped like ``query``. + + Raises: + RuntimeError: The cuDNN Frontend package is unavailable. + """ + execute = _build_cudnn_fp8_sdpa( + query.device, + torch.cuda.current_stream(query.device).cuda_stream, + tuple(query.shape), + query.stride(), + tuple(key.shape), + key.stride(), + tuple(value.shape), + value.stride(), + ) + return execute(query, key, value) diff --git a/flashdreams/flashdreams/accelerated/multi_head_attention/cudnn/torch_sdpa.py b/flashdreams/flashdreams/accelerated/multi_head_attention/cudnn/torch_sdpa.py new file mode 100644 index 000000000..c66b759ef --- /dev/null +++ b/flashdreams/flashdreams/accelerated/multi_head_attention/cudnn/torch_sdpa.py @@ -0,0 +1,35 @@ +# 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. + +"""PyTorch scaled-dot-product attention forced to cuDNN.""" + +import torch +import torch.nn.functional as F +from torch import Tensor + + +def torch_cudnn_sdpa(query: Tensor, key: Tensor, value: Tensor) -> Tensor: + """Apply PyTorch scaled-dot-product attention with the cuDNN backend. + + Args: + query: Queries in ``[B, H, L, D]`` layout. + key: Keys in ``[B, H, S, D]`` layout. + value: Values in ``[B, H, S, D]`` layout. + + Returns: + Attention output shaped like ``query``. + """ + with torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.CUDNN_ATTENTION): + return F.scaled_dot_product_attention(query, key, value) diff --git a/flashdreams/flashdreams/accelerated/multi_head_attention/optimized.py b/flashdreams/flashdreams/accelerated/multi_head_attention/optimized.py new file mode 100644 index 000000000..f7f426f1d --- /dev/null +++ b/flashdreams/flashdreams/accelerated/multi_head_attention/optimized.py @@ -0,0 +1,1110 @@ +# 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. + +"""Optimized multi-head attention with fused projections and selectable SDPA.""" + +from __future__ import annotations + +import math +from abc import abstractmethod +from collections.abc import Callable +from dataclasses import dataclass +from enum import Enum + +import torch +from torch import Tensor, nn + +from flashdreams.accelerated.common.non_persistent_linear import ( + NonPersistentLinear, +) +from flashdreams.accelerated.multi_head_attention import ( + AttentionConfig, + AttentionType, + MultiHeadAttention, + QKNormScope, + RoPEScope, + RoPEStyle, +) +from flashdreams.accelerated.multi_head_attention.cudnn import ( + native_cudnn_fp8_sdpa, + torch_cudnn_sdpa, +) +from flashdreams.accelerated.multi_head_attention.triton import ( + flash_attention_2, + flash_attention_2_tma, + is_tma_flash_attention_supported, +) +from flashdreams.accelerated.quantization.linear import ( + QuantizedNonPersistentLinear, + WeightGranularity, +) +from flashdreams.accelerated.quantization.quantizer import ( + DTYPE_MAX, + Granularity, + quantize, +) +from flashdreams.core.attention import BlockKVCache +from flashdreams.core.attention.rope_kernel import apply_rotary_pos_emb + + +class SDPABackend(str, Enum): + """Scaled-dot-product attention implementation.""" + + CUDNN = "cudnn" + """Use Torch cuDNN for FP16/BF16 and native cuDNN Frontend for FP8.""" + + FA2 = "fa2" + """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.""" + + +@dataclass(frozen=True, slots=True) +class QuantizationOption: + """Attention quantization policy.""" + + projection: torch.dtype | None = None + """Q/K/V projection dtype; ``None`` preserves native precision.""" + + quantized_sdpa: bool = False + """Use unscaled FP8 e4m3 Q/K/V in scaled-dot-product attention. + + This directly casts Q, K, and V to FP8 e4m3 before calling the configured + SDPA backend and stores the K/V cache in that dtype. The cuDNN path requires + ``nvidia-cudnn-frontend``. FA2 also casts the + softmax probabilities used by ``P @ V`` to FP8 e4m3. This is not + SageAttention3 quantization: it uses no accuracy-preserving quantization + scheme or scaling. As noted by the SageAttention3 paper, this simple approach + can make attention inaccurate, so output accuracy is not guaranteed. + """ + + def __post_init__(self) -> None: + """Validate the projection quantization dtype.""" + if self.projection is not None and self.projection not in DTYPE_MAX: + raise ValueError( + f"unsupported projection quantization dtype: {self.projection}" + ) + if not isinstance(self.quantized_sdpa, bool): + raise TypeError( + f"quantized_sdpa must be a bool; got {self.quantized_sdpa!r}" + ) + + +@dataclass(frozen=True, slots=True) +class OptimizedImplConfig: + """Optimized attention implementation policy.""" + + qkv_fusion_option: QKVFusionOption = QKVFusionOption.FULL + """Projection fusion policy.""" + + sdpa_backend: SDPABackend = SDPABackend.CUDNN + """Scaled-dot-product attention implementation.""" + + use_tma: bool = True + """Prefer TMA FlashAttention2 when the device and tensors support it.""" + + quantization: QuantizationOption = QuantizationOption() + """Attention quantization policy.""" + + def __post_init__(self) -> None: + """Validate optimized implementation policy values.""" + if not isinstance(self.qkv_fusion_option, QKVFusionOption): + raise TypeError( + "qkv_fusion_option must be a QKVFusionOption; " + f"got {self.qkv_fusion_option!r}" + ) + if not isinstance(self.quantization, QuantizationOption): + raise TypeError( + f"quantization must be a QuantizationOption; got {self.quantization!r}" + ) + if not isinstance(self.sdpa_backend, SDPABackend): + raise TypeError( + f"sdpa_backend must be an SDPABackend; got {self.sdpa_backend!r}" + ) + if not isinstance(self.use_tma, bool): + raise TypeError(f"use_tma must be a bool; got {self.use_tma!r}") + + +class OptimizedHultiHeadAttention(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 before PyTorch Q/K + normalization, the shared RoPE kernel, and standard cache updates. + 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.""" + + optimized_impl_config: OptimizedImplConfig + """Projection and attention backend policy.""" + + fused_qkv: NonPersistentLinear | None + """Nonpersistent full-fusion Q/K/V projection.""" + + fused_kv: NonPersistentLinear | None + """Nonpersistent fused K/V projection.""" + + quantized_query_projection: QuantizedNonPersistentLinear | None + """Nonpersistent quantized query projection.""" + + quantized_key_projection: QuantizedNonPersistentLinear | None + """Nonpersistent quantized key projection.""" + + quantized_value_projection: QuantizedNonPersistentLinear | None + """Nonpersistent quantized value projection.""" + + _validated_cuda_device_index: int | None + """CUDA device whose compute capability has passed validation.""" + + def __init__( + self, + attention_type: AttentionType, + attention_config: AttentionConfig, + optimized_impl_config: OptimizedImplConfig, + ) -> None: + """Initialize shared optimized attention policies. + + Args: + attention_type: Select self-attention or static cross-attention. + attention_config: Shared geometry, normalization, and rotary policy. + optimized_impl_config: Projection and attention backend policy. + + Raises: + TypeError: ``optimized_impl_config`` has the wrong type. + ValueError: A fusion or head dimension is unsupported. + """ + super().__init__(attention_type, attention_config) + if not isinstance(optimized_impl_config, OptimizedImplConfig): + raise TypeError( + "optimized_impl_config must be an OptimizedImplConfig; " + f"got {optimized_impl_config!r}" + ) + qkv_fusion_option = optimized_impl_config.qkv_fusion_option + if ( + qkv_fusion_option is QKVFusionOption.FULL + and self.attention_config.query_dim != self.attention_config.context_dim + ): + raise ValueError( + "full QKV fusion requires query_dim to equal context_dim; " + f"got {self.attention_config.query_dim} and {self.attention_config.context_dim}" + ) + if not ( + 16 <= self.attention_config.head_dim <= 256 + and self.attention_config.head_dim & (self.attention_config.head_dim - 1) + == 0 + ): + raise ValueError( + "accelerated attention requires a power-of-two head_dim in [16, 256]; " + f"got {self.attention_config.head_dim}" + ) + + self.optimized_impl_config = optimized_impl_config + self.qkv_fusion_option = optimized_impl_config.qkv_fusion_option + self.sdpa_backend = optimized_impl_config.sdpa_backend + self.use_tma = optimized_impl_config.use_tma + self._validated_cuda_device_index = None + + # ---------------------- Initialization ---------------------- # + + def _initialize_derived_weights(self) -> None: + """Build fused execution modules after concrete checkpoint fields exist. + + Concrete implementations call this after assigning their checkpoint + fields. The logical accessors are valid before fused projections or load + hooks read any projection or normalization module. + """ + self.fused_qkv = None + self.fused_kv = None + self.quantized_query_projection = None + self.quantized_key_projection = None + self.quantized_value_projection = None + self._refresh_derived_weights() + self.register_load_state_dict_post_hook(self._refresh_derived_weights) + + @staticmethod + def _new_quantized_projection( + weight: Tensor, + bias: Tensor | None, + dtype: torch.dtype, + ) -> QuantizedNonPersistentLinear: + """Build a per-output-channel quantized projection from checkpoint tensors.""" + return QuantizedNonPersistentLinear( + weight.detach().contiguous(), + None if bias is None else bias.detach(), + WeightGranularity.PER_OUT_CHANNEL, + dtype, + ) + + @torch.no_grad() + def _refresh_derived_weights(self, *args: object) -> None: + """Rebuild fused projection modules from checkpoint parameters.""" + del args + self.fused_qkv = None + self.fused_kv = None + self.quantized_query_projection = None + self.quantized_key_projection = None + self.quantized_value_projection = None + + projection_dtype = self.optimized_impl_config.quantization.projection + if projection_dtype is not None: + self.quantized_query_projection = self._new_quantized_projection( + self.query_projection.weight, + self.query_projection.bias, + projection_dtype, + ) + if self.qkv_fusion_option is QKVFusionOption.NONE: + self.quantized_key_projection = self._new_quantized_projection( + self.key_projection.weight, + self.key_projection.bias, + projection_dtype, + ) + self.quantized_value_projection = self._new_quantized_projection( + self.value_projection.weight, + self.value_projection.bias, + projection_dtype, + ) + + if self.qkv_fusion_option is QKVFusionOption.FULL: + fused_weight = ( + torch.cat( + ( + self.query_projection.weight, + self.key_projection.weight, + self.value_projection.weight, + ), + dim=0, + ) + .detach() + .contiguous() + ) + fused_bias = 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 + ) + fused_bias = torch.cat( + ( + self.query_projection.bias, + self.key_projection.bias, + self.value_projection.bias, + ), + dim=0, + ).detach() + fused_kv_weight = fused_weight[self.attention_config.inner_dim :] + fused_kv_bias = ( + None + if fused_bias is None + else fused_bias[self.attention_config.inner_dim :] + ) + if projection_dtype is None: + self.fused_qkv = NonPersistentLinear(fused_weight, fused_bias) + self.fused_kv = NonPersistentLinear(fused_kv_weight, fused_kv_bias) + else: + self.fused_qkv = self._new_quantized_projection( + fused_weight, fused_bias, projection_dtype + ) + self.fused_kv = self._new_quantized_projection( + fused_kv_weight, fused_kv_bias, projection_dtype + ) + elif self.qkv_fusion_option is QKVFusionOption.FUSE_KV: + fused_weight = ( + torch.cat( + (self.key_projection.weight, self.value_projection.weight), dim=0 + ) + .detach() + .contiguous() + ) + fused_bias = None + if self.key_projection.bias is not None: + assert self.value_projection.bias is not None + fused_bias = torch.cat( + (self.key_projection.bias, self.value_projection.bias), dim=0 + ).detach() + if projection_dtype is None: + self.fused_kv = NonPersistentLinear(fused_weight, fused_bias) + else: + self.fused_kv = self._new_quantized_projection( + fused_weight, fused_bias, projection_dtype + ) + + def _apply( + self, + fn: Callable[[Tensor], Tensor], + recurse: bool = True, + ) -> OptimizedHultiHeadAttention: + """Transform parameters and rebuild fused modules 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 fused projection modules refreshed. + """ + # Move/cast canonical parameters first, then regenerate fused projections. + 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 for native or quantized SDPA. + + 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: FP16 or BF16 activation dtype. Quantized SDPA stores the cache + in FP8 e4m3 instead. + + Returns: + Block cache with K/V storage shaped + ``[B, sink_size + window_size, H, D]``. + + """ + # 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.attention_config.n_heads, + self.attention_config.head_dim, + ) + self._validate_cuda_device(device) + return 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=( + torch.float8_e4m3fn + if self.optimized_impl_config.quantization.quantized_sdpa + else dtype + ), + ) + + @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 rotations shaped ``[S, 1, 1, D]``. + Applied only by before-cache RoPE; after-cache RoPE stores + unrotated keys. Ignored when ``rope_config`` is ``None``. + + 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 ( + self.attention_config.rope_config is not None + and self.attention_config.rope_config.scope is RoPEScope.BEFORE_KV_CACHE + and rope_freqs is not None + ): + # Position affects key directions used in Q·K; values are never rotated. + key = self._apply_rope(key, rope_freqs) + + if self.optimized_impl_config.quantization.quantized_sdpa: + key = key.to(torch.float8_e4m3fn) + value = value.to(torch.float8_e4m3fn) + + 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 positional data. Before-cache RoPE expects the + current chunk. After-cache RoPE expects positions relative to + the visible cache. Ignored when ``rope_config`` is ``None``. + + Returns: + Output-projected tokens with the same shape and dtype as ``x``. + """ + query_rope_freqs, key_rope_freqs = self._slice_rope_freqs(rope_freqs, kv_cache) + if self.attention_type is AttentionType.SELF_ATTENTION: + query = self._update_kv_and_compute_query(x, kv_cache, query_rope_freqs) + else: + query = self._compute_query(x, query_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. + key = kv_cache.cached_k() + if ( + self.attention_config.rope_config is not None + and self.attention_config.rope_config.scope is RoPEScope.AFTER_KV_CACHE + and key_rope_freqs is not None + ): + # The shared RoPE kernel is in-place; keep cache storage unrotated so + # rolling positions can be applied again on the next attention call. + key = self._apply_rope(key.to(x.dtype, copy=True), key_rope_freqs) + value = kv_cache.cached_v() + if self.optimized_impl_config.quantization.quantized_sdpa: + query = query.to(torch.float8_e4m3fn) + key = key.to(torch.float8_e4m3fn) + value = value.to(torch.float8_e4m3fn) + output = self._attention( + query, + key, + value, + output_dtype=x.dtype, + ) + sequence_length = x.shape[-2] + output = output.reshape(-1, sequence_length, self.attention_config.inner_dim) + output = self._project_output(output) + return output.reshape( + x.shape[:-2] + (sequence_length, self.attention_config.query_dim) + ) + + # ------------------------------------------------------------ # + # Private Method # + # ------------------------------------------------------------ # + + # ------------------ Core Attention Methods ------------------ # + + def _slice_rope_freqs( + self, + rope_freqs: Tensor | None, + kv_cache: BlockKVCache, + ) -> tuple[Tensor | None, Tensor | None]: + """Select query and key rotations for the configured cache scope. + + Args: + rope_freqs: Current-chunk or cache-relative rotation angles. + kv_cache: Cache whose visible and current write ranges select angles. + + Returns: + Query and visible-key rotation slices, or two ``None`` values when + rotary embeddings are disabled. + """ + if self.attention_config.rope_config is None or rope_freqs is None: + return None, None + if self.attention_config.rope_config.scope is RoPEScope.BEFORE_KV_CACHE: + return rope_freqs, rope_freqs + + write_end = kv_cache.write_end + write_start = write_end - kv_cache.chunk_size + return rope_freqs[write_start:write_end], rope_freqs[: kv_cache.size] + + def _compute_query( + self, + query: Tensor, + rope_freqs: Tensor | None, + prequantized_input: tuple[Tensor, Tensor] | None = None, + ) -> Tensor: + """Project, normalize, and optionally rotate query tokens. + + Args: + query: Full-precision query tokens. + rope_freqs: Optional query rotations. + prequantized_input: Quantized ``query`` and its scale, or ``None`` + to quantize a separate query projection in this call. + + Returns: + Processed query tensor in token-major head layout. + """ + if self.attention_config.rope_config is None: + rope_freqs = None + query = self._project_query(query, prequantized_input) + if self.attention_config.rope_config is not None and rope_freqs is not None: + query = self._apply_rope(query, rope_freqs) + return 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.attention_config.rope_config is None: + rope_freqs = None + 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.attention_config.query_dim) + query, key, value = self._project_qkv(x_flat) + query = self._apply_qk_norm(query, self.query_norm) + key = self._apply_qk_norm(key, self.key_norm) + if rope_freqs is not None: + query = self._apply_rope(query, rope_freqs) + if ( + self.attention_config.rope_config is not None + and self.attention_config.rope_config.scope + is RoPEScope.BEFORE_KV_CACHE + ): + key = self._apply_rope(key, rope_freqs) + if self.optimized_impl_config.quantization.quantized_sdpa: + key = key.to(torch.float8_e4m3fn) + value = value.to(torch.float8_e4m3fn) + kv_cache.update(key, value) + return query + + assert self.attention_config.context_dim is not None + self._validate_tokens(x, self.attention_config.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]}" + ) + + prequantized_input = None + if self.optimized_impl_config.quantization.projection is not None: + prequantized_input = self._prequantize_projection_input(x) + query = self._compute_query(x, rope_freqs, prequantized_input) + key, value = self._project_kv(x, prequantized_input) + if ( + self.attention_config.rope_config is not None + and self.attention_config.rope_config.scope is RoPEScope.BEFORE_KV_CACHE + and rope_freqs is not None + ): + key = self._apply_rope(key, rope_freqs) + if self.optimized_impl_config.quantization.quantized_sdpa: + key = key.to(torch.float8_e4m3fn) + value = value.to(torch.float8_e4m3fn) + kv_cache.update(key, value) + return query + + def _attention( + self, + query: Tensor, + key: Tensor, + value: Tensor, + *, + output_dtype: torch.dtype | None = None, + ) -> 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]``. + output_dtype: Output storage dtype; ``None`` uses ``query.dtype``. + + 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 change only shape/stride metadata. + query = query.transpose(1, 2) + key = key.transpose(1, 2) + value = value.transpose(1, 2) + + # PyTorch's public dispatcher rejects FP8 inputs, so use a cuDNN + # Frontend FP8 graph for e4m3 attention. + if query.dtype is torch.float8_e4m3fn: + output = native_cudnn_fp8_sdpa(query, key, value) + else: + output = torch_cudnn_sdpa(query, key, value) + + # Restore the module-wide ``[B, L, H, D]`` contract for head merging. + output = output.transpose(1, 2) + return output if output_dtype is None else output.to(output_dtype) + + attention = ( + flash_attention_2_tma + if self.use_tma and is_tma_flash_attention_supported(query, key, value) + else flash_attention_2 + ) + if output_dtype is None: + return attention(query, key, value) + return attention(query, key, value, output_dtype=output_dtype) + + def _apply_rope(self, x: Tensor, rope_freqs: Tensor) -> Tensor: + """Apply the shared RoPE kernel to token-major head features. + + Args: + x: Projected Q or K tensor shaped ``[B, L, H, D]``. + rope_freqs: Rotation angles shaped ``[L, 1, 1, D]``. + + Returns: + In-place rotated tensor with the same shape and dtype as ``x``. + """ + rope_config = self.attention_config.rope_config + if rope_config is None: + return x + return apply_rotary_pos_emb( + x, + rope_freqs, + interleaved=rope_config.style is RoPEStyle.INTERLEAVED, + inplace=True, + ) + + # ------------------------ Validation ------------------------ # + + def _validate_cuda_device(self, device: torch.device | str) -> None: + """Validate a CUDA device once before it enters the hot path. + + Args: + device: Device used by attention inputs and cache storage. + + Raises: + RuntimeError: The CUDA device predates Hopper. + """ + device = torch.device(device) + if device.type != "cuda": + return + device_index = ( + torch.cuda.current_device() if device.index is None else device.index + ) + if self._validated_cuda_device_index == device_index: + return + if torch.cuda.get_device_capability(device_index)[0] < 9: + raise RuntimeError( + "OptimizedHultiHeadAttention requires compute capability 9.0 or newer" + ) + self._validated_cuda_device_index = device_index + + 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( + "OptimizedHultiHeadAttention requires CUDA FP16 or BF16 inputs" + ) + self._validate_cuda_device(x.device) + + 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( + "OptimizedHultiHeadAttention 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.attention_config.n_heads, + self.attention_config.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("Optimized 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.optimized_impl_config.quantization.quantized_sdpa + 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.attention_config.query_dim: + raise ValueError( + f"x feature width must equal query_dim={self.attention_config.query_dim}; " + f"got {x.shape[-1]}" + ) + + # The accelerated path accepts native FP16/BF16 CUDA inputs. + if not x.is_cuda or x.dtype not in (torch.float16, torch.bfloat16): + raise RuntimeError( + "OptimizedHultiHeadAttention requires CUDA FP16 or BF16 inputs" + ) + self._validate_cuda_device(x.device) + + # 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( + "OptimizedHultiHeadAttention 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.attention_config.n_heads, + self.attention_config.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, dtype, and dense layout so the standard + # cache update preserves the layout required by the attention backend. + if kv_cache._v.shape != kv_cache._k.shape: + raise ValueError("Optimized 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.optimized_impl_config.quantization.quantized_sdpa + 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.attention_config.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 self.attention_config.rope_config is not None and rope_freqs is not None: + # RoPE coefficients cover this ``L``-token chunk and broadcast across + # flattened batches and ``H`` heads inside the shared kernel. + expected_rope_shape = (x.shape[-2], 1, 1, self.attention_config.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 _prequantize_projection_input(self, x: Tensor) -> tuple[Tensor, Tensor]: + """Quantize one source tensor for reuse across separate projections. + + Args: + x: Full-precision projection input. + + Returns: + Quantized activations and their per-slice scale. + """ + projection_dtype = self.optimized_impl_config.quantization.projection + assert projection_dtype is not None + return quantize(x, projection_dtype, Granularity.SLICE, axis=-1) + + def _apply_qk_norm(self, x: Tensor, norm: nn.Module) -> Tensor: + """Apply PyTorch Q/K normalization with the configured feature scope. + + Args: + x: Projected query or key shaped ``[B, L, H, D]``. + norm: PyTorch normalization module or identity. + + Returns: + Normalized tensor with the same shape as ``x``. + """ + if self.attention_config.qk_norm_scope is QKNormScope.INNER: + return norm(x.flatten(-2)).reshape(x.shape) + return norm(x) + + def _project_query( + self, + query: Tensor, + prequantized_input: tuple[Tensor, Tensor] | None = None, + ) -> Tensor: + """Project and normalize queries in token-major head layout.""" + self._validate_tokens(query, self.attention_config.query_dim, "query") + sequence_length = query.shape[-2] + if self.optimized_impl_config.quantization.projection is None: + projected_query = self.query_projection(query) + else: + if self.quantized_query_projection is None: + raise RuntimeError("quantized query projection is not initialized") + if prequantized_input is None: + prequantized_input = self._prequantize_projection_input(query) + quantized_query, query_scale = prequantized_input + projected_query = self.quantized_query_projection( + quantized_query, + query_scale, + out_dtype=query.dtype, + ) + projected_query = projected_query.reshape( + -1, + sequence_length, + self.attention_config.n_heads, + self.attention_config.head_dim, + ) + return self._apply_qk_norm(projected_query, self.query_norm) + + def _project_kv( + self, + context: Tensor, + prequantized_input: tuple[Tensor, Tensor] | None = None, + ) -> tuple[Tensor, Tensor]: + """Project and normalize context keys in token-major head layout.""" + assert self.attention_config.context_dim is not None + self._validate_tokens(context, self.attention_config.context_dim, "context") + sequence_length = context.shape[-2] + head_shape = ( + -1, + sequence_length, + self.attention_config.n_heads, + self.attention_config.head_dim, + ) + if self.qkv_fusion_option is QKVFusionOption.NONE: + if self.optimized_impl_config.quantization.projection is None: + projected_key = self.key_projection(context) + projected_value = self.value_projection(context) + else: + if ( + self.quantized_key_projection is None + or self.quantized_value_projection is None + ): + raise RuntimeError("quantized K/V projections are not initialized") + if prequantized_input is None: + prequantized_input = self._prequantize_projection_input(context) + quantized_context, context_scale = prequantized_input + projected_key = self.quantized_key_projection( + quantized_context, + context_scale, + out_dtype=context.dtype, + ) + projected_value = self.quantized_value_projection( + quantized_context, + context_scale, + out_dtype=context.dtype, + ) + key = projected_key.reshape(head_shape) + value = projected_value.reshape(head_shape) + else: + if self.fused_kv is None: + raise RuntimeError("fused K/V projection is not initialized") + if self.optimized_impl_config.quantization.projection is None: + projected_kv = self.fused_kv(context) + else: + if not isinstance(self.fused_kv, QuantizedNonPersistentLinear): + raise RuntimeError( + "quantized fused K/V projection is not initialized" + ) + if prequantized_input is None: + projected_kv = self.fused_kv( + context, + Granularity.SLICE, + out_dtype=context.dtype, + ) + else: + quantized_context, context_scale = prequantized_input + projected_kv = self.fused_kv( + quantized_context, + context_scale, + out_dtype=context.dtype, + ) + projected_kv = projected_kv.reshape( + -1, + sequence_length, + 2, + self.attention_config.n_heads, + self.attention_config.head_dim, + ) + key, value = projected_kv.unbind(dim=2) + return self._apply_qk_norm(key, self.key_norm), value + + def _project_qkv(self, x: Tensor) -> tuple[Tensor, Tensor, Tensor]: + """Project Q/K/V with one fused GEMM.""" + if self.fused_qkv is None: + raise RuntimeError("fused QKV projection is not initialized") + if self.optimized_impl_config.quantization.projection is None: + qkv = self.fused_qkv(x) + else: + if not isinstance(self.fused_qkv, QuantizedNonPersistentLinear): + raise RuntimeError("quantized fused QKV projection is not initialized") + qkv = self.fused_qkv(x, Granularity.SLICE, out_dtype=x.dtype) + qkv = qkv.reshape( + -1, + x.shape[-2], + 3, + self.attention_config.n_heads, + self.attention_config.head_dim, + ) + query, key, value = qkv.unbind(dim=2) + return query, key, value + + def _project_output(self, x: Tensor) -> Tensor: + """Apply the FP16/BF16 output projection.""" + return self.output_projection(x) + + +__all__ = [ + "QKVFusionOption", + "QuantizationOption", + "SDPABackend", + "OptimizedImplConfig", + "OptimizedHultiHeadAttention", + "flash_attention_2", + "flash_attention_2_tma", + "is_tma_flash_attention_supported", +] diff --git a/flashdreams/flashdreams/accelerated/multi_head_attention/torch.py b/flashdreams/flashdreams/accelerated/multi_head_attention/torch.py new file mode 100644 index 000000000..fede44e9a --- /dev/null +++ b/flashdreams/flashdreams/accelerated/multi_head_attention/torch.py @@ -0,0 +1,543 @@ +# 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. + +"""PyTorch self- and cross-attention over block K/V caches.""" + +from __future__ import annotations + +from abc import abstractmethod + +import torch +import torch.nn.functional as F +from torch import Tensor, nn + +from flashdreams.accelerated.multi_head_attention import ( + AttentionConfig, + AttentionType, + MultiHeadAttention, + QKNormScope, + RoPEScope, + RoPEStyle, +) +from flashdreams.core.attention import BlockKVCache + + +class TorchMultiHeadAttention(MultiHeadAttention[BlockKVCache]): + """PyTorch reference for self- and cross-attention with a block cache. + + This implementation owns the cache lifecycle dispatched by :meth:`forward`. + Self-attention writes each current chunk to a prepared rolling cache; + cross-attention reuses K/V materialized once by :meth:`compute_kv`. + + Shape descriptions use ``B`` for the product of all leading batch or grouping + dimensions, ``L`` for query tokens, ``S`` for visible cached context, ``H`` + for heads, and ``D`` for head features. Projections collapse leading token + dimensions into the cache's single ``B`` axis; query outputs restore their + original leading layout. Query and context layouts may therefore differ as + long as their flattened batch sizes agree. + + Native PyTorch SDPA supplies the portable reference backend. Concrete + subclasses own the projection and normalization modules and can override + :meth:`_attention` while retaining the RoPE and cache contracts. + """ + + @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.""" + + def __init__( + self, + attention_type: AttentionType, + attention_config: AttentionConfig, + ) -> None: + """Initialize shared reference-attention policies. + + Args: + attention_type: Select self-attention or static cross-attention. + attention_config: Shared geometry, normalization, and rotary policy. + """ + super().__init__(attention_type, attention_config) + + 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 native-precision rolling K/V cache. + + The returned cache is empty. For every chunk, call + :meth:`BlockKVCache.before_update`, invoke self-attention, and call + :meth:`BlockKVCache.after_update` with the same chunk index. + + Args: + batch_size: Product ``B`` of the input's leading dimensions. + chunk_size: Exact number of current tokens ``L`` 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: Data type used by K/V storage. + + Returns: + Block cache with K/V storage shaped + ``[B, sink_size + window_size, H, D]``. + """ + # BlockKVCache rolls sequence dimension 1 while preserving independent + # batch and head axes for SDPA. + cache_shape = ( + batch_size, + sink_size + window_size, + self.attention_config.n_heads, + self.attention_config.head_dim, + ) + return 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=dtype, + ) + + def compute_kv( + self, + context: Tensor, + rope_freqs: Tensor | None = None, + ) -> BlockKVCache: + """Project static context into a finalized, reusable K/V cache. + + Args: + context: Static key/value source, shape + ``[..., S, context_dim]``. Leading dimensions flatten into ``B``. + rope_freqs: Optional key rotations with shape ``[S, 1, 1, D]``. + Applied only by before-cache RoPE; after-cache RoPE stores + unrotated keys. Ignored when ``rope_config`` is ``None``. + + Returns: + Static cache with K/V shape ``[B, S, H, D]``, ready for repeated + :meth:`forward` calls without lifecycle bookkeeping. + + Raises: + ValueError: Context width or RoPE geometry is incompatible with the + configured attention geometry. + """ + key, value = self._project_kv(context) + if ( + self.attention_config.rope_config is not None + and self.attention_config.rope_config.scope is RoPEScope.BEFORE_KV_CACHE + and rope_freqs is not None + ): + key = self._apply_rope(key, rope_freqs) + # ``from_tensor`` completes its one write internally, so static context + # can be queried immediately and never enters the rolling lifecycle. + return BlockKVCache.from_tensor(key, value, seq_dim=1) + + def forward( + self, + x: Tensor, + kv_cache: BlockKVCache, + rope_freqs: Tensor | None = None, + ) -> Tensor: + """Apply self- or cross-attention using the configured cache lifecycle. + + Args: + x: Query tokens, shape ``[..., L, query_dim]``. + kv_cache: Prepared rolling cache for self-attention or precomputed + static cache for cross-attention. + rope_freqs: Optional positional data. Before-cache RoPE expects the + current chunk. After-cache RoPE expects positions relative to + the visible cache. Ignored when ``rope_config`` is ``None``. + + Returns: + Output-projected tokens with the same shape as ``x``. + """ + query_rope_freqs, key_rope_freqs = self._slice_rope_freqs(rope_freqs, kv_cache) + if self.attention_type is AttentionType.SELF_ATTENTION: + kv_cache = self._update_kv(x, kv_cache, key_rope_freqs) + return self._query_kv(x, kv_cache, query_rope_freqs, key_rope_freqs) + + def _slice_rope_freqs( + self, + rope_freqs: Tensor | None, + kv_cache: BlockKVCache, + ) -> tuple[Tensor | None, Tensor | None]: + """Select query and key rotations for the configured cache scope. + + Args: + rope_freqs: Current-chunk or cache-relative rotation angles. + kv_cache: Cache whose visible and current write ranges select angles. + + Returns: + Query and visible-key rotation slices, or two ``None`` values when + rotary embeddings are disabled. + """ + if self.attention_config.rope_config is None or rope_freqs is None: + return None, None + if self.attention_config.rope_config.scope is RoPEScope.BEFORE_KV_CACHE: + return rope_freqs, rope_freqs + + write_end = kv_cache.write_end + write_start = write_end - kv_cache.chunk_size + return rope_freqs[write_start:write_end], rope_freqs[: kv_cache.size] + + def _update_kv( + self, + context: Tensor, + kv_cache: BlockKVCache, + rope_freqs: Tensor | None = None, + ) -> BlockKVCache: + """Project and write one context chunk to a prepared rolling cache. + + Args: + context: Current key/value source chunk, shape + ``[..., L, context_dim]`` where ``L`` equals + ``kv_cache.chunk_size``. + kv_cache: Rolling cache after + :meth:`BlockKVCache.before_update` for the current chunk. + rope_freqs: Optional key rotations. Before-cache RoPE expects + shape ``[L, 1, 1, D]``; after-cache RoPE stores unrotated keys. + + Returns: + The same cache with the current K/V chunk written and visible to + attention. The caller still owns final bookkeeping. + + Raises: + ValueError: Token width, RoPE geometry, cache geometry, or chunk + length is incompatible with this module. + RuntimeError: Cache device or dtype differs from the projected + tensors, or its update lifecycle is inactive. + """ + key, value = self._project_kv(context) + if ( + self.attention_config.rope_config is not None + and self.attention_config.rope_config.scope is RoPEScope.BEFORE_KV_CACHE + and rope_freqs is not None + ): + key = self._apply_rope(key, rope_freqs) + self._validate_cache(kv_cache, key, updating=True) + kv_cache.update(key, value) + # Leave the lifecycle open so a following query sees the current chunk; + # the caller closes it with ``after_update`` after attention completes. + return kv_cache + + def _query_kv( + self, + query: Tensor, + kv_cache: BlockKVCache, + rope_freqs: Tensor | None = None, + key_rope_freqs: Tensor | None = None, + ) -> Tensor: + """Query visible K/V without mutating cache storage or bookkeeping. + + Args: + query: Query tokens, shape ``[..., L, query_dim]``. Its flattened + leading size must equal the cache batch size ``B``. + kv_cache: Static or rolling cache exposing K/V as ``[B, S, H, D]``. + rope_freqs: Optional query rotations with shape ``[L, 1, 1, D]``; + ``None`` leaves queries unrotated. + key_rope_freqs: Optional visible-cache rotations with shape + ``[S, 1, 1, D]`` for after-cache RoPE. + + Returns: + Output-projected tokens with the same leading dimensions and shape + ``[..., L, query_dim]`` as ``query``. + + Raises: + ValueError: Query width, RoPE geometry, or cache geometry is + incompatible with this module. + RuntimeError: Cache device or dtype differs from projected queries. + """ + # Projection folds every leading query dimension into the cache's single + # batch axis; retain the public layout for the final reshape. + batch_shape = query.shape[:-2] + query = self._project_query(query) + if self.attention_config.rope_config is not None and rope_freqs is not None: + query = self._apply_rope(query, rope_freqs) + self._validate_cache(kv_cache, query) + key = kv_cache.cached_k() + if ( + self.attention_config.rope_config is not None + and self.attention_config.rope_config.scope is RoPEScope.AFTER_KV_CACHE + and key_rope_freqs is not None + ): + key = self._apply_rope(key, key_rope_freqs) + value = kv_cache.cached_v() + output = self._attention(query, key, value) + output = self._output_projection(output) + # ``output`` is ``[B, L, query_dim]`` after head concatenation; restore + # the exact leading query geometry captured at the module boundary. + return output.reshape(batch_shape + output.shape[-2:]) + + def _validate_tokens(self, x: Tensor, feature_dim: int, name: str) -> None: + """Validate a public token tensor's rank and feature width. + + Args: + x: Query or context tokens with expected shape ``[..., L, C]``. + feature_dim: Required trailing feature width ``C``. + name: Argument name included in validation errors. + + Raises: + ValueError: ``x`` has fewer than two dimensions or the wrong trailing + feature width. + """ + 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]}" + ) + + def _validate_cache( + self, + kv_cache: BlockKVCache, + x: Tensor, + *, + updating: bool = False, + ) -> None: + """Validate cache geometry, placement, and optional write lifecycle. + + Args: + kv_cache: Block cache expected to expose ``[B, S, H, D]`` K/V. + x: Projected queries or keys with shape ``[B, L, H, D]``. + updating: Also require an active update and ``L`` equal to the cache + chunk size. + + Raises: + ValueError: Cache layout, batch/head geometry, or update length is + incompatible with ``x``. + RuntimeError: Cache device or dtype differs from ``x``, or + ``updating`` is true outside an active update. + """ + # The reference cache ABI is token-major ``[B, S, H, D]`` with a + # dynamically visible ``S``; only the other axes are fixed here. + if kv_cache.seq_dim != 1 or kv_cache._k.ndim != 4: + raise ValueError("K/V cache must have shape [B, S, H, D] with seq_dim=1") + if kv_cache._v.shape != kv_cache._k.shape: + raise ValueError("K/V cache tensors must have identical shapes") + expected = ( + x.shape[0], + self.attention_config.n_heads, + self.attention_config.head_dim, + ) + actual = ( + kv_cache._k.shape[0], + kv_cache._k.shape[2], + kv_cache._k.shape[3], + ) + if actual != expected: + raise ValueError( + "cache batch, head, and feature dimensions must equal " + f"{expected}; got {actual}" + ) + 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") + if kv_cache._k.dtype != x.dtype or kv_cache._v.dtype != x.dtype: + raise RuntimeError("K/V cache tensors must match the input dtype") + if not updating: + return + if kv_cache._curr_chunk_idx is None: + raise RuntimeError("call kv_cache.before_update() before attention") + if x.shape[1] != kv_cache.chunk_size: + raise ValueError( + "context sequence length must equal cache chunk_size=" + f"{kv_cache.chunk_size}; got {x.shape[1]}" + ) + + def _project_query(self, query: Tensor) -> Tensor: + """Project and normalize queries as ``[B, L, H, D]``. + + Args: + query: Query tokens with shape ``[..., L, query_dim]``. + + Returns: + Projected queries with all leading dimensions flattened into ``B``. + """ + self._validate_tokens(query, self.attention_config.query_dim, "query") + sequence_length = query.shape[-2] + # Split ``H * D`` while collapsing arbitrary batch/group dimensions into + # the one batch axis shared with BlockKVCache. + query = self.query_projection(query).reshape( + -1, + sequence_length, + self.attention_config.n_heads, + self.attention_config.head_dim, + ) + if self.attention_config.qk_norm_scope is QKNormScope.INNER: + # INNER normalization sees concatenated heads; HEAD and NONE operate + # directly on trailing ``D`` through RMSNorm or Identity. + return self.query_norm(query.flatten(-2)).reshape(query.shape) + return self.query_norm(query) + + def _project_kv(self, context: Tensor) -> tuple[Tensor, Tensor]: + """Project context into K/V shaped ``[B, S, H, D]``. + + Args: + context: Key/value source with shape ``[..., S, context_dim]``. + + Returns: + Normalized keys and unnormalized values with flattened batch size + ``B`` and independent ``H`` and ``D`` axes. + """ + assert self.attention_config.context_dim is not None + self._validate_tokens(context, self.attention_config.context_dim, "context") + sequence_length = context.shape[-2] + # K/V share the cache ABI even when context leading dimensions differ + # from the query layout used later. + head_shape = ( + -1, + sequence_length, + self.attention_config.n_heads, + self.attention_config.head_dim, + ) + key = self.key_projection(context).reshape(head_shape) + value = self.value_projection(context).reshape(head_shape) + if self.attention_config.qk_norm_scope is QKNormScope.INNER: + # Values are never Q/K-normalized; only keys follow the configured + # head or concatenated-inner normalization policy. + key = self.key_norm(key.flatten(-2)).reshape(key.shape) + else: + key = self.key_norm(key) + return key, value + + def _apply_rope(self, x: Tensor, rope_freqs: Tensor) -> Tensor: + """Apply the configured RoPE pairing to projected queries or keys. + + Args: + x: Projected tensor with shape ``[B, L, H, D]``. + rope_freqs: Rotation angles with exact shape ``[L, 1, 1, D]``. + + Returns: + Rotated tensor with the same shape, device, and dtype as ``x``. + + Raises: + ValueError: ``D`` is odd or ``rope_freqs`` has incompatible geometry. + """ + if self.attention_config.rope_config is None: + return x + if x.shape[-1] % 2 != 0: + raise ValueError(f"RoPE requires an even head_dim; got {x.shape[-1]}") + + # RoPE lookup shape is ``[L, 1, 1, D]`` for an input shaped + # ``[..., L, H, D]``. + 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)}" + ) + + # Broadcast positions over leading dimensions and heads: + # ``[L, 1, 1, D] -> [..., L, 1, D]``. + freqs = rope_freqs[:, 0, 0, :].reshape( + (1,) * (x.ndim - 3) + (x.shape[-3], 1, x.shape[-1]) + ) + + # Materialize rotation coefficients in activation precision so the + # elementwise rotation neither promotes projected tensors nor cache data. + cos_freqs = torch.cos(freqs).to(dtype=x.dtype) + sin_freqs = torch.sin(freqs).to(dtype=x.dtype) + if self.attention_config.rope_config.style is RoPEStyle.INTERLEAVED: + # Rotate adjacent feature pairs; shape stays ``[..., L, H, D]``. + rotated = torch.stack((-x[..., 1::2], x[..., 0::2]), dim=-1).flatten(-2) + else: + # Rotate matching half-split features; shape stays ``[..., L, H, D]``. + first, second = x.chunk(2, dim=-1) + rotated = torch.cat((-second, first), dim=-1) + + # Apply the elementwise complex rotation: ``[..., L, H, D]``. + return x * cos_freqs + rotated * sin_freqs + + def _attention(self, query: Tensor, key: Tensor, value: Tensor) -> Tensor: + """Apply non-causal scaled dot-product attention over visible K/V. + + Args: + query: Projected queries with shape ``[B, L, H, D]``. + key: Visible cached keys with shape ``[B, S, H, D]``. + value: Visible cached values with shape ``[B, S, H, D]``. + + Returns: + Per-head attention output with shape ``[B, L, H, D]``. + """ + # Move heads before tokens for SDPA: + # Q ``[..., L, H, D] -> [..., H, L, D]`` and + # K/V ``[..., S, H, D] -> [..., H, S, D]``. + query_heads = query.transpose(-3, -2) + key_heads = key.transpose(-3, -2) + value_heads = value.transpose(-3, -2) + + # Let PyTorch dispatch the available SDPA backend so the reference works + # on CPU and CUDA. Cache visibility defines the allowed context, while + # zero dropout and a non-causal mask make inference deterministic. + output = F.scaled_dot_product_attention( + query_heads, + key_heads, + value_heads, + dropout_p=0.0, + is_causal=False, + ) + + # Restore token-major layout: ``[..., H, L, D] -> [..., L, H, D]``. + return output.transpose(-3, -2) + + def _output_projection(self, x: Tensor) -> Tensor: + """Concatenate attention heads and project back to query features. + + Args: + x: Per-head attention output with shape ``[B, L, H, D]``. + + Returns: + Projected tokens with shape ``[B, L, query_dim]``. + """ + # Concatenate heads: ``[..., L, H, D] -> [..., L, H * D]``. + x = x.flatten(-2) + + # Project attention features: ``[..., L, H * D] -> [..., L, query_dim]``. + return self.output_projection(x) + + +__all__ = ["TorchMultiHeadAttention"] diff --git a/flashdreams/flashdreams/accelerated/multi_head_attention/triton/__init__.py b/flashdreams/flashdreams/accelerated/multi_head_attention/triton/__init__.py new file mode 100644 index 000000000..059a9c1d1 --- /dev/null +++ b/flashdreams/flashdreams/accelerated/multi_head_attention/triton/__init__.py @@ -0,0 +1,30 @@ +# 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 FlashAttention2 kernels.""" + +from flashdreams.accelerated.multi_head_attention.triton.flash_attention_2_kernel import ( + flash_attention_2, +) +from flashdreams.accelerated.multi_head_attention.triton.flash_attention_2_tma_kernel import ( + flash_attention_2_tma, + is_tma_flash_attention_supported, +) + +__all__ = [ + "flash_attention_2", + "flash_attention_2_tma", + "is_tma_flash_attention_supported", +] diff --git a/flashdreams/flashdreams/accelerated/multi_head_attention/triton/flash_attention_2_kernel.py b/flashdreams/flashdreams/accelerated/multi_head_attention/triton/flash_attention_2_kernel.py new file mode 100644 index 000000000..5575c4351 --- /dev/null +++ b/flashdreams/flashdreams/accelerated/multi_head_attention/triton/flash_attention_2_kernel.py @@ -0,0 +1,374 @@ +# 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. + +"""Pointer-based Triton FlashAttention2 for projected attention tensors.""" + +from __future__ import annotations + +import math + +import torch +import triton +import triton.language as tl +from torch import Tensor + +_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 load-pipeline resource use.""" + + +def _prune_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 and wide heads do + not compile configurations whose padded work or accumulator footprint cannot + be competitive. + + Args: + configs: Candidate autotuning configurations. + named_args: Runtime arguments containing ``query_length`` and + ``key_length``. + **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"] + head_dim = meta["HEAD_DIM"] + assert isinstance(query_length, int) + assert isinstance(key_length, 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. + 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, 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 + ] + + +# Cache the winning tile by logical geometry and sequence strides. Pointer values +# and the numeric softmax scale do not change scheduling, +# so they intentionally do not create new autotuning entries. + + +@triton.autotune( + configs=_ATTENTION_CONFIGS, + key=[ + "num_heads", + "query_length", + "key_length", + "query_stride_l", + "key_stride_s", + "value_stride_s", + "HEAD_DIM", + ], + prune_configs_by={"early_config_prune": _prune_attention_configs}, + cache_results=True, +) +@triton.jit +def _flash_attention_2_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, + scale, + HEAD_DIM: tl.constexpr, + QUANTIZED_SDPA: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + """Apply tiled non-causal FlashAttention2 with pointer 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. Each program loads one ``[BLOCK_M, D]`` query tile, streams + every ``[BLOCK_N, D]`` K/V tile, and writes the matching output tile. + + 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 attention heads. + query_length: Logical query-token count ``L``. + key_length: Logical key/value-token count ``S``. + scale: Multiplier applied to QK scores before softmax. + HEAD_DIM: Compile-time head width ``D``. + QUANTIZED_SDPA: Whether Q/K/V and P use FP8 e4m3. + BLOCK_M: Compile-time number of query rows owned by one program. + BLOCK_N: Compile-time number of key/value rows loaded per iteration. + """ + query_block = tl.program_id(0) + batch_head = tl.program_id(1) + batch = batch_head // num_heads + head = batch_head % num_heads + + 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_offsets = query_block * BLOCK_M + tl.arange(0, BLOCK_M) + feature_offsets = tl.arange(0, HEAD_DIM) + query_mask = query_offsets < query_length + query = tl.load( + query_base + + query_offsets[:, None] * query_stride_l + + feature_offsets[None, :] * query_stride_d, + mask=query_mask[:, None], + other=0.0, + ) + + 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) + qk_scale = scale.to(tl.float32) * 1.4426950408889634 + + for key_start in tl.range(0, key_length, BLOCK_N): + key_offsets = key_start + tl.arange(0, BLOCK_N) + key_mask = key_offsets < key_length + key = tl.load( + key_base + + key_offsets[:, None] * key_stride_s + + feature_offsets[None, :] * key_stride_d, + mask=key_mask[:, None], + other=0.0, + ) + scores = tl.dot(query, tl.trans(key)) * qk_scale + scores = tl.where(key_mask[None, :], scores, -float("inf")) + + 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) + + value = tl.load( + value_base + + key_offsets[:, None] * value_stride_s + + feature_offsets[None, :] * value_stride_d, + mask=key_mask[:, None], + other=0.0, + ) + accumulator *= correction[:, None] + if QUANTIZED_SDPA: + probabilities = probabilities.to(tl.float8e4nv) + else: + probabilities = probabilities.to(value.dtype) + accumulator = tl.dot(probabilities, value, accumulator) + row_max = next_row_max + + output = accumulator / denominator[:, None] + tl.store( + output_base + + query_offsets[:, None] * output_stride_l + + feature_offsets[None, :] * output_stride_d, + output, + mask=query_mask[:, None], + ) + + +def flash_attention_2( + query: Tensor, + key: Tensor, + value: Tensor, + *, + scale: float | None = None, + output_dtype: torch.dtype | None = None, +) -> Tensor: + """Apply non-causal pointer-based FlashAttention2 to Q/K/V tensors. + + Compute ``softmax(scale * Q @ K.T) @ V`` independently for every + batch/head plane without dropout or materializing the complete score matrix. + Empty batch, head, or query axes return an empty output, but the key/value + sequence axis must be positive. + + Args: + query: CUDA FP16, BF16, or FP8 e4m3 query tensor with shape + ``[B, L, H, D]``. + 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)``. + output_dtype: Output storage dtype; ``None`` uses ``query.dtype``. + + Returns: + Attention result with shape ``[B, L, H, D]`` on the query device and + in ``output_dtype``. + + Raises: + ValueError: Q/K/V shapes are incompatible or contain an empty key axis. + RuntimeError: Placement, dtype, head geometry, or strides do not satisfy + the pointer-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 query.is_cuda or not key.is_cuda or not value.is_cuda: + raise RuntimeError("FlashAttention2 requires CUDA tensors") + if query.device != key.device or query.device != value.device: + raise RuntimeError("query, key, and value must occupy the same CUDA device") + if query.dtype != key.dtype or query.dtype != value.dtype: + raise RuntimeError("query, key, and value must have the same dtype") + if query.dtype not in ( + torch.float16, + torch.bfloat16, + torch.float8_e4m3fn, + ): + raise RuntimeError("FlashAttention2 requires FP16, BF16, or FP8 e4m3 tensors") + if not (16 <= head_dim <= 256 and head_dim & (head_dim - 1) == 0): + raise RuntimeError( + "FlashAttention2 requires a power-of-two head_dim in [16, 256]" + ) + if any(stride <= 0 for x in (query, key, value) for stride in x.stride()): + raise RuntimeError("FlashAttention2 requires positive tensor strides") + + if output_dtype is None: + output_dtype = query.dtype + if output_dtype not in (torch.float16, torch.bfloat16, torch.float8_e4m3fn): + raise RuntimeError("FlashAttention2 requires an FP16, BF16, or FP8 e4m3 output") + output = torch.empty(query.shape, device=query.device, dtype=output_dtype) + if batch_size == 0 or num_heads == 0 or query_length == 0: + return output + + 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), + ) + + 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, + ) + + _flash_attention_2_kernel[grid]( + query, + key, + value, + output, + *query_strides, + *key_strides, + *value_strides, + *output_strides, + num_heads, + query_length, + key_length, + 1.0 / math.sqrt(head_dim) if scale is None else scale, + HEAD_DIM=head_dim, + QUANTIZED_SDPA=query.dtype is torch.float8_e4m3fn, + ) + return output + + +__all__ = ["flash_attention_2"] diff --git a/flashdreams/flashdreams/accelerated/multi_head_attention/triton/flash_attention_2_tma_kernel.py b/flashdreams/flashdreams/accelerated/multi_head_attention/triton/flash_attention_2_tma_kernel.py new file mode 100644 index 000000000..0814c9af6 --- /dev/null +++ b/flashdreams/flashdreams/accelerated/multi_head_attention/triton/flash_attention_2_tma_kernel.py @@ -0,0 +1,494 @@ +# 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 +import triton +import triton.language as tl +from torch import Tensor + + +def _descriptor_layout_supported(x: Tensor) -> bool: + """Return whether ``x`` satisfies TMA tensor-descriptor stride rules. + + 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() + 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 satisfy the TMA kernel contract. + """ + 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 _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 and wide heads do + not compile configurations whose padded work or accumulator footprint cannot + be competitive. + + Args: + configs: Candidate autotuning configurations. + named_args: Runtime arguments containing ``query_length`` and + ``key_length``. + **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"] + head_dim = meta["HEAD_DIM"] + assert isinstance(query_length, int) + assert isinstance(key_length, 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. + 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, 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 + ] + + +# Cache the winning tile by logical geometry and sequence strides. 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", + "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, + scale, + HEAD_DIM: tl.constexpr, + QUANTIZED_SDPA: 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``. + scale: Multiplier applied to QK scores before softmax. + HEAD_DIM: Compile-time head width ``D``. + QUANTIZED_SDPA: Whether Q/K/V and P use FP8 e4m3. + 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] + if QUANTIZED_SDPA: + probabilities = probabilities.to(tl.float8e4nv) + else: + probabilities = probabilities.to(value.dtype) + accumulator = tl.dot(probabilities, 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 flash_attention_2_tma( + query: Tensor, + key: Tensor, + value: Tensor, + *, + scale: float | None = None, + output_dtype: torch.dtype | 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. Empty + batch, head, or query axes return an empty output, but + the key/value sequence axis must be positive. + + Args: + query: CUDA FP16, BF16, or FP8 e4m3 query tensor with shape + ``[B, L, H, D]``. + 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)``. + output_dtype: Output storage dtype; ``None`` uses ``query.dtype``. + + Returns: + Attention result with shape ``[B, L, H, D]`` on the query device and in + ``output_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 e4m3 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. + if output_dtype is None: + output_dtype = query.dtype + if output_dtype not in (torch.float16, torch.bfloat16, torch.float8_e4m3fn): + raise RuntimeError( + "TMA FlashAttention2 requires an FP16, BF16, or FP8 e4m3 output" + ) + output = torch.empty( + query.shape, + device=query.device, + dtype=output_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, + # 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, + ) + + _flash_attention_2_tma_kernel[grid]( + query, + key, + value, + output, + *query_strides, + *key_strides, + *value_strides, + *output_strides, + num_heads, + query_length, + key_length, + 1.0 / math.sqrt(head_dim) if scale is None else scale, + HEAD_DIM=head_dim, + QUANTIZED_SDPA=query.dtype is torch.float8_e4m3fn, + ) + return output + + +__all__ = ["flash_attention_2_tma"] diff --git a/flashdreams/flashdreams/accelerated/quantization/linear.py b/flashdreams/flashdreams/accelerated/quantization/linear.py new file mode 100644 index 000000000..33838d99e --- /dev/null +++ b/flashdreams/flashdreams/accelerated/quantization/linear.py @@ -0,0 +1,306 @@ +# 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. + +"""Quantized nonpersistent linear transformation for accelerated inference.""" + +from enum import Enum +from typing import overload + +import torch +from torch import Tensor + +from flashdreams.accelerated.common.non_persistent_linear import ( + NonPersistentLinear, +) +from flashdreams.accelerated.quantization.quantizer import ( + Granularity, + dequantize, + quantize, +) + + +class WeightGranularity(str, Enum): + """Scale granularity for quantized linear weights.""" + + PER_OUT_CHANNEL = "per_out_channel" + """Use one scale for every output channel.""" + + TENSOR = "tensor" + """Use one scale for the complete weight tensor.""" + + +class QuantizedNonPersistentLinear(NonPersistentLinear): + """Apply a nonpersistent linear transformation with quantized weights. + + The layer quantizes and stores ``weight [O, I]`` during construction, + where ``I`` is ``in_features`` and ``O`` is ``out_features``. At inference, + it accepts activations ``x [..., I]`` and returns ``output [..., O]``. The + leading activation dimensions may represent any combination of batch, + sequence, or spatial dimensions. + + Examples: + Construct a layer and let it quantize full-precision activations using + one scale per ``I``-element slice:: + + import torch + + from flashdreams.accelerated.quantization.linear import ( + QuantizedNonPersistentLinear, + WeightGranularity, + ) + from flashdreams.accelerated.quantization.quantizer import ( + Granularity, + quantize, + ) + + weight = torch.randn(32, 64, device="cuda", dtype=torch.float16) + bias = torch.randn(32, device="cuda", dtype=torch.float16) + layer = QuantizedNonPersistentLinear( + weight, + bias, + WeightGranularity.PER_OUT_CHANNEL, + torch.float8_e4m3fn, + ) + x = torch.randn(2, 8, 64, device="cuda", dtype=torch.float16) + output = layer(x, Granularity.SLICE) + assert output.shape == (2, 8, 32) + + Reuse prequantized activations and their scale to avoid quantizing + ``x`` inside each call:: + + quantized_x, x_scale = quantize( + x, + layer.dtype, + Granularity.SLICE, + axis=-1, + ) + output = layer(quantized_x, x_scale, out_dtype=torch.float16) + assert output.shape == (2, 8, 32) + """ + + dtype: torch.dtype + """Quantized activation format required by prequantized inputs.""" + + weight_scale: Tensor + """FP32 weight scale shaped ``[O, 1]`` or ``[1, 1]``.""" + + def __init__( + self, + weight: Tensor, + bias: Tensor | None, + granularity: WeightGranularity, + dtype: torch.dtype, + ) -> None: + """Initialize the transformation from existing tensors. + + Args: + weight: Projection weight shaped ``[out_features, in_features]``. + bias: Optional projection bias shaped ``[out_features]``. + granularity: Scale granularity used to quantize ``weight``. + dtype: Quantized activation format. + + Raises: + ValueError: ``granularity`` or ``dtype`` is unsupported. + """ + # Map the public weight modes onto ``quantize`` with ``axis=-1``: + # ``[O, I] -> scale [O, 1]`` per output or ``scale [1, 1]`` per tensor. + if granularity is WeightGranularity.PER_OUT_CHANNEL: + quantizer_granularity = Granularity.SLICE + elif granularity is WeightGranularity.TENSOR: + quantizer_granularity = Granularity.TENSOR + else: + raise ValueError(f"unsupported weight granularity: {granularity}") + + # CUDA FP8 GEMM pairs E5M2 activations with an E4M3 weight operand. + # All other formats use the same dtype for ``x`` and ``weight``. + weight_dtype = torch.float8_e4m3fn if dtype is torch.float8_e5m2 else dtype + quantized_weight, weight_scale = quantize( + weight, + weight_dtype, + quantizer_granularity, + axis=-1, + ) + # Keep derived ``weight [O, I]`` and its scale out of ``state_dict``; + # callers recreate both from the source weight when building the layer. + super().__init__(quantized_weight.contiguous(), bias) + self.dtype = dtype + self.register_buffer( + "weight_scale", weight_scale.contiguous(), persistent=False + ) + + @overload + def forward( + self, + x: Tensor, + scale_or_granularity: Tensor, + out_dtype: torch.dtype = torch.float16, + ) -> Tensor: ... + + @overload + def forward( + self, + x: Tensor, + scale_or_granularity: Granularity, + out_dtype: torch.dtype = torch.float16, + ) -> Tensor: ... + + def forward( + self, + x: Tensor, + scale_or_granularity: Tensor | Granularity, + out_dtype: torch.dtype = torch.float16, + ) -> Tensor: + """Apply the quantized linear transformation. + + Args: + x: Activations shaped ``[..., in_features]``. When a scale tensor + is supplied, these must already use the layer's quantized + ``dtype``; when a granularity is supplied, these are + full-precision activations quantized inside this call. + scale_or_granularity: For prequantized ``x``, an FP32 tensorwise + scale shaped ``[1, ..., 1]`` or slice scale shaped + ``[..., 1]``. For full-precision ``x``, the granularity used + to produce one of those scale shapes. + out_dtype: Data type of the projected activations. Defaults to + ``torch.float16``. + + Returns: + Projected activations shaped ``[..., out_features]``. + + Raises: + ValueError: ``x`` or its scale does not match the layer contract. + """ + if x.ndim == 0 or x.shape[-1] != self.in_features: + raise ValueError( + f"expected input last dim {self.in_features}, got {tuple(x.shape)}" + ) + + if isinstance(scale_or_granularity, Tensor): + # Prequantized path: preserve ``x [..., I]`` and its existing + # tensorwise ``scale [1, ..., 1]`` or slice ``scale [..., 1]``. + if x.dtype is not self.dtype: + raise ValueError( + f"expected quantized input dtype {self.dtype}, got {x.dtype}" + ) + self._validate_scale(x, scale_or_granularity) + quantized, scale = x, scale_or_granularity + elif isinstance(scale_or_granularity, Granularity): + # Dynamic path: quantization keeps ``x [..., I]`` unchanged in + # shape and reduces either every dimension or only ``I`` for scale. + if x.numel() == 0: + return torch.empty( + (*x.shape[:-1], self.out_features), + device=x.device, + dtype=out_dtype, + ) + quantized, scale = quantize( + x, + self.dtype, + scale_or_granularity, + axis=-1, + ) + else: + raise ValueError( + "scale_or_granularity must be a scale tensor or Granularity" + ) + + return self._forward_quantized(quantized, scale, out_dtype) + + @staticmethod + def _validate_scale(x: Tensor, scale: Tensor) -> None: + """Validate a tensorwise ``[1, ..., 1]`` or slice ``[..., 1]`` scale.""" + tensor_shape = (1,) * x.ndim + slice_shape = (*x.shape[:-1], 1) + if scale.shape not in (tensor_shape, slice_shape): + raise ValueError( + f"expected scale shape {tensor_shape} or {slice_shape}, " + f"got {tuple(scale.shape)}" + ) + if scale.dtype is not torch.float32: + raise ValueError(f"expected FP32 scale, got {scale.dtype}") + if scale.device != x.device: + raise ValueError(f"expected scale on device {x.device}, got {scale.device}") + + def _forward_quantized( + self, x: Tensor, scale: Tensor, out_dtype: torch.dtype + ) -> Tensor: + """Project validated ``x [..., I]`` into ``output [..., O]``.""" + if x.numel() == 0: + return torch.empty( + (*x.shape[:-1], self.out_features), + device=x.device, + dtype=out_dtype, + ) + + # Collapse all leading dimensions into GEMM rows: + # ``x [..., I] -> input_2d [R, I]``, where ``R = prod(x.shape[:-1])``. + input_2d = x.reshape(-1, self.in_features).contiguous() + # A slice scale ``[..., 1]`` follows the same collapse to ``[R, 1]``; + # a tensor scale ``[1, ..., 1]`` becomes the scalar matrix ``[1, 1]``. + input_scale = scale.reshape(-1, 1).contiguous() + # Transpose weight scales from quantizer layout ``[O, 1]`` or + # ``[1, 1]`` to GEMM output-column layout ``[1, O]`` or ``[1, 1]``. + weight_scale = self.weight_scale.T.contiguous() + + if self.dtype is torch.int8: + # Multiply ``[R, I] @ [I, O] -> int32 [R, O]``, then broadcast + # activation scales down rows and weight scales across columns. + input_rows = input_2d.shape[0] + if input_rows <= 16: + # CUDA ``_int_mm`` requires more than 16 rows. Zero-padding the + # integer GEMM and slicing before dequantization preserves the + # original activation scales and output exactly. + input_2d = torch.cat( + ( + input_2d, + input_2d.new_zeros(17 - input_rows, self.in_features), + ) + ) + output = dequantize( + torch._int_mm(input_2d, self.weight.T)[:input_rows], + input_scale, + weight_scale, + dtype=out_dtype, + ) + else: + # ``_scaled_mm`` accepts either two ``[1, 1]`` scales or rowwise + # ``[R, 1]`` and ``[1, O]`` scales. Expand only the scalar side for + # mixed granularities; repeated values preserve its tensor scale. + if input_scale.numel() == 1 and weight_scale.numel() != 1: + input_scale = input_scale.expand(input_2d.shape[0], 1).contiguous() + elif weight_scale.numel() == 1 and input_scale.numel() != 1: + weight_scale = weight_scale.expand(1, self.out_features).contiguous() + + rowwise_scaling = input_scale.numel() != 1 or weight_scale.numel() != 1 + # The CUDA rowwise path produces reliable high-precision ``[R, O]`` + # through BF16; cast after GEMM when the caller requests otherwise. + scaled_out_dtype = torch.bfloat16 if rowwise_scaling else out_dtype + # Multiply ``input_2d [R, I]`` by ``weight.T [I, O]`` and fuse both + # dequantization scales into the resulting ``output [R, O]``. + output = torch._scaled_mm( + input_2d, + self.weight.T, + input_scale, + weight_scale, + out_dtype=scaled_out_dtype, + ) + if output.dtype is not out_dtype: + output = output.to(out_dtype) + + # Broadcast ``bias [O]`` over the ``R`` rows, then restore the original + # leading dimensions: ``output [R, O] -> [..., O]``. + if self.bias is not None: + output = output + self.bias.to(device=output.device, dtype=out_dtype) + return output.reshape(*x.shape[:-1], self.out_features) diff --git a/flashdreams/flashdreams/accelerated/quantization/quantizer.py b/flashdreams/flashdreams/accelerated/quantization/quantizer.py new file mode 100644 index 000000000..d6a48321e --- /dev/null +++ b/flashdreams/flashdreams/accelerated/quantization/quantizer.py @@ -0,0 +1,123 @@ +# 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. + +"""Quantization granularity and tensor conversion interfaces.""" + +from enum import Enum + +import torch +from torch import Tensor + +from flashdreams.accelerated.quantization.quantizer_kernel import ( + dequantize_triton, + quantize_triton, +) + +DTYPE_MAX: dict[torch.dtype, float] = { + torch.float8_e4m3fn: torch.finfo(torch.float8_e4m3fn).max, + torch.float8_e5m2: torch.finfo(torch.float8_e5m2).max, + torch.int8: torch.iinfo(torch.int8).max, +} +"""Largest finite positive value for each supported quantized dtype.""" + + +class Granularity(str, Enum): + """Scale granularity for tensor quantization.""" + + SLICE = "slice" + """Use one scale per slice selected by ``axis`` in :func:`quantize`.""" + + TENSOR = "tensor" + """Use one scale for the complete tensor.""" + + +def quantize( + original: Tensor, + format: torch.dtype, + granularity: Granularity, + axis: int = -1, + use_triton: bool = True, +) -> tuple[Tensor, Tensor]: + """Quantize a tensor using the requested format and scale granularity. + + Args: + original: Tensor to quantize. + format: Data type of the quantized tensor. + granularity: Scope over which scales are computed. ``SLICE`` reduces + along ``axis``; ``TENSOR`` reduces over the complete tensor. + axis: Dimension over which ``SLICE`` computes the maximum value. For a + two-dimensional tensor, ``axis=1`` produces one scale per row and + ``axis=0`` produces one scale per column. Defaults to ``-1``. + use_triton: Use Triton for CUDA tensors. CPU tensors retain the Torch + implementation. Defaults to ``True``. + + Returns: + Quantized tensor and its FP32 scale tensor. The scale has the same + number of dimensions as ``original`` and retains reduced dimensions + with size one. + + Raises: + ValueError: ``format`` or ``granularity`` is unsupported. + """ + if format not in DTYPE_MAX: + raise ValueError(f"unsupported quantization format: {format}") + + if granularity is Granularity.TENSOR: + reduction_axis: int | tuple[int, ...] = tuple(range(original.ndim)) + elif granularity is Granularity.SLICE: + reduction_axis = axis + else: + raise ValueError(f"unsupported quantization granularity: {granularity}") + + if use_triton and original.is_cuda: + return quantize_triton(original, format, granularity, axis) + + original_float = original.detach().to(torch.float32) + max_abs = original_float.abs().amax(dim=reduction_axis, keepdim=True) + scale = (max_abs / DTYPE_MAX[format]).clamp_min(torch.finfo(torch.float32).tiny) + quantized = (original_float / scale).clamp(-DTYPE_MAX[format], DTYPE_MAX[format]) + if not format.is_floating_point: + quantized = quantized.round() + return quantized.to(format), scale + + +def dequantize( + quantized: Tensor, + *scales: Tensor, + dtype: torch.dtype = torch.float16, + use_triton: bool = True, +) -> Tensor: + """Dequantize a tensor using its scale tensors. + + Args: + quantized: Tensor to dequantize. + scales: Scale tensors used to dequantize ``quantized``. + dtype: Data type of the dequantized tensor. Defaults to ``torch.float16``. + use_triton: Use Triton for CUDA tensors. CPU tensors retain the Torch + implementation. Defaults to ``True``. + + Returns: + Dequantized tensor in ``dtype`` after applying every scale in order. + Without scales, casts ``quantized`` directly to ``dtype``. + """ + if use_triton and quantized.is_cuda: + return dequantize_triton(quantized, *scales, dtype=dtype) + if not scales: + return quantized.to(dtype) + + dequantized = quantized.to(scales[0].dtype) + for scale in scales: + dequantized = dequantized * scale + return dequantized.to(dtype) diff --git a/flashdreams/flashdreams/accelerated/quantization/quantizer_kernel.py b/flashdreams/flashdreams/accelerated/quantization/quantizer_kernel.py new file mode 100644 index 000000000..d871bc6fe --- /dev/null +++ b/flashdreams/flashdreams/accelerated/quantization/quantizer_kernel.py @@ -0,0 +1,376 @@ +# 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 kernels for quantizing and dequantizing tensors.""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl +from torch import Tensor +from triton.language.extra import libdevice + +_SUPPORTED_FORMATS = ( + torch.float8_e4m3fn, + torch.float8_e5m2, + torch.int8, +) + +_ELEMENT_BLOCK_SIZE = 1024 +"""Number of contiguous elements processed by elementwise programs.""" + +_MAX_REDUCTION_BLOCK_SIZE = 16384 +"""Largest power-of-two reduction tile kept in one Triton program.""" + +_FLOAT32_TINY = tl.constexpr(torch.finfo(torch.float32).tiny) +"""Smallest positive normal FP32 value used for zero-valued scales.""" + + +@triton.jit +def _quantize_values( + values, scale, max_value: tl.constexpr, round_values: tl.constexpr +): + scaled = tl.maximum( + tl.minimum(libdevice.div_rn(values, scale), max_value), -max_value + ) + if round_values: + scaled = libdevice.rint(scaled) + return scaled + + +@triton.jit +def _quantize_slices_kernel( + original_ptr, + quantized_ptr, + scale_ptr, + axis_size, + inner_size, + max_value: tl.constexpr, + round_values: tl.constexpr, + block_size: tl.constexpr, +): + group_index = tl.program_id(0) + inner_index = group_index % inner_size + outer_index = group_index // inner_size + group_start = outer_index * axis_size * inner_size + inner_index + offsets = tl.arange(0, block_size) + + max_abs = 0.0 + for start in range(0, axis_size, block_size): + axis_offsets = start + offsets + values = tl.load( + original_ptr + group_start + axis_offsets * inner_size, + mask=axis_offsets < axis_size, + other=0.0, + ).to(tl.float32) + max_abs = tl.maximum(max_abs, tl.max(tl.abs(values), axis=0)) + + scale = tl.maximum(max_abs / max_value, _FLOAT32_TINY) + tl.store(scale_ptr + group_index, scale) + + for start in range(0, axis_size, block_size): + axis_offsets = start + offsets + mask = axis_offsets < axis_size + pointers = original_ptr + group_start + axis_offsets * inner_size + values = tl.load(pointers, mask=mask, other=0.0).to(tl.float32) + quantized = _quantize_values(values, scale, max_value, round_values) + tl.store( + quantized_ptr + group_start + axis_offsets * inner_size, + quantized, + mask, + ) + + +@triton.jit +def _partial_max_kernel( + original_ptr, partial_max_ptr, element_count, block_size: tl.constexpr +): + offsets = tl.program_id(0) * block_size + tl.arange(0, block_size) + values = tl.load( + original_ptr + offsets, + mask=offsets < element_count, + other=0.0, + ).to(tl.float32) + tl.store(partial_max_ptr + tl.program_id(0), tl.max(tl.abs(values), axis=0)) + + +@triton.jit +def _tensor_scale_kernel( + partial_max_ptr, + scale_ptr, + partial_count, + max_value: tl.constexpr, + block_size: tl.constexpr, +): + offsets = tl.arange(0, block_size) + max_abs = 0.0 + for start in range(0, partial_count, block_size): + partial_offsets = start + offsets + values = tl.load( + partial_max_ptr + partial_offsets, + mask=partial_offsets < partial_count, + other=0.0, + ) + max_abs = tl.maximum(max_abs, tl.max(values, axis=0)) + tl.store(scale_ptr, tl.maximum(max_abs / max_value, _FLOAT32_TINY)) + + +@triton.jit +def _quantize_tensor_kernel( + original_ptr, + quantized_ptr, + scale_ptr, + element_count, + max_value: tl.constexpr, + round_values: tl.constexpr, + block_size: tl.constexpr, +): + offsets = tl.program_id(0) * block_size + tl.arange(0, block_size) + mask = offsets < element_count + values = tl.load(original_ptr + offsets, mask=mask, other=0.0).to(tl.float32) + scale = tl.load(scale_ptr) + quantized = _quantize_values(values, scale, max_value, round_values) + tl.store(quantized_ptr + offsets, quantized, mask) + + +@triton.jit +def _cast_kernel(input_ptr, output_ptr, element_count, block_size: tl.constexpr): + offsets = tl.program_id(0) * block_size + tl.arange(0, block_size) + mask = offsets < element_count + tl.store(output_ptr + offsets, tl.load(input_ptr + offsets, mask=mask), mask) + + +@triton.jit +def _multiply_broadcast_kernel( + input_ptr, + scale_ptr, + output_ptr, + element_count, + shape: tl.constexpr, + scale_strides: tl.constexpr, + block_size: tl.constexpr, +): + offsets = tl.program_id(0) * block_size + tl.arange(0, block_size) + remaining = offsets + scale_offsets = tl.zeros([block_size], tl.int64) + for dimension in tl.static_range(len(shape) - 1, -1, -1): + coordinate = remaining % shape[dimension] + remaining = remaining // shape[dimension] + scale_offsets += coordinate * scale_strides[dimension] + + mask = offsets < element_count + values = tl.load(input_ptr + offsets, mask=mask, other=0.0) + scales = tl.load(scale_ptr + scale_offsets, mask=mask, other=0.0) + tl.store(output_ptr + offsets, values * scales, mask) + + +def _reduction_block_size(element_count: int) -> int: + """Return a bounded power-of-two reduction tile size.""" + return min(triton.next_power_of_2(element_count), _MAX_REDUCTION_BLOCK_SIZE) + + +def _num_warps(block_size: int) -> int: + """Return enough warps for the selected reduction tile.""" + return 8 if block_size >= 2048 else 4 + + +def _format_max(format: torch.dtype) -> float: + """Return the largest finite value for a supported quantized format.""" + if format not in _SUPPORTED_FORMATS: + raise ValueError(f"unsupported quantization format: {format}") + if format.is_floating_point: + return torch.finfo(format).max + return torch.iinfo(format).max + + +def _normalize_axis(axis: int, ndim: int) -> int: + """Normalize ``axis`` while retaining Torch's scalar-axis behavior.""" + if ndim == 0: + if axis not in (-1, 0): + raise IndexError( + "Dimension out of range (expected to be in range of [-1, 0], " + f"but got {axis})" + ) + return 0 + if axis < -ndim or axis >= ndim: + raise IndexError( + "Dimension out of range (expected to be in range of " + f"[{-ndim}, {ndim - 1}], but got {axis})" + ) + return axis % ndim + + +def quantize_triton( + original: Tensor, + format: torch.dtype, + granularity: str, + axis: int = -1, +) -> tuple[Tensor, Tensor]: + """Quantize a CUDA tensor with tensorwise or axis-slice scales. + + Args: + original: CUDA tensor to quantize. + format: Quantized output data type. + granularity: ``"tensor"`` or ``"slice"`` scale granularity. + axis: Reduction dimension for slice granularity. Defaults to ``-1``. + + Returns: + Quantized tensor and its FP32 scale tensor with reduced dimensions kept. + + Raises: + ValueError: The input is not a CUDA tensor or an option is unsupported. + IndexError: ``axis`` is invalid or selects an empty reduction dimension. + """ + if not original.is_cuda: + raise ValueError("Triton quantization requires a CUDA tensor") + + max_value = _format_max(format) + original = original.detach().contiguous() + quantized = torch.empty_like(original, dtype=format) + + if granularity == "tensor": + if original.numel() == 0: + raise IndexError("amax(): Expected reduction dim to have non-zero size") + scale = torch.empty( + (1,) * original.ndim, + device=original.device, + dtype=torch.float32, + ) + partial_count = triton.cdiv(original.numel(), _ELEMENT_BLOCK_SIZE) + partial_max = torch.empty( + partial_count, + device=original.device, + dtype=torch.float32, + ) + _partial_max_kernel[(partial_count,)]( + original, + partial_max, + original.numel(), + block_size=_ELEMENT_BLOCK_SIZE, + ) + reduction_block_size = _reduction_block_size(partial_count) + _tensor_scale_kernel[(1,)]( + partial_max, + scale, + partial_count, + max_value=max_value, + block_size=reduction_block_size, + num_warps=_num_warps(reduction_block_size), + ) + _quantize_tensor_kernel[(partial_count,)]( + original, + quantized, + scale, + original.numel(), + max_value=max_value, + round_values=format is torch.int8, + block_size=_ELEMENT_BLOCK_SIZE, + ) + return quantized, scale + + if granularity != "slice": + raise ValueError(f"unsupported quantization granularity: {granularity}") + + normalized_axis = _normalize_axis(axis, original.ndim) + axis_size = original.shape[normalized_axis] if original.ndim else 1 + if axis_size == 0: + raise IndexError( + f"amax(): Expected reduction dim {axis} to have non-zero size." + ) + scale_shape = list(original.shape) + if scale_shape: + scale_shape[normalized_axis] = 1 + scale = torch.empty(scale_shape, device=original.device, dtype=torch.float32) + if original.numel() == 0: + return quantized, scale + + inner_size = original.stride(normalized_axis) if original.ndim else 1 + block_size = _reduction_block_size(axis_size) + _quantize_slices_kernel[(scale.numel(),)]( + original, + quantized, + scale, + axis_size, + inner_size, + max_value=max_value, + round_values=format is torch.int8, + block_size=block_size, + num_warps=_num_warps(block_size), + ) + return quantized, scale + + +def dequantize_triton( + quantized: Tensor, + *scales: Tensor, + dtype: torch.dtype = torch.float16, +) -> Tensor: + """Dequantize a CUDA tensor while applying broadcastable scales in order. + + Args: + quantized: CUDA tensor to dequantize. + scales: Scale tensors broadcastable to ``quantized``. + dtype: Output data type. Defaults to ``torch.float16``. + + Returns: + Dequantized tensor in ``dtype``. + + Raises: + ValueError: A tensor is not on the same CUDA device as ``quantized``. + RuntimeError: A scale cannot broadcast to ``quantized``. + """ + if not quantized.is_cuda: + raise ValueError("Triton dequantization requires a CUDA tensor") + if any(scale.device != quantized.device for scale in scales): + raise ValueError("dequantization scales must share the quantized tensor device") + if quantized.numel() == 0: + dequantized = quantized.to(scales[0].dtype) if scales else quantized + for scale in scales: + dequantized = dequantized * scale + return dequantized.to(dtype) + + shape = tuple(quantized.shape) + element_count = quantized.numel() + grid = (triton.cdiv(element_count, _ELEMENT_BLOCK_SIZE),) + current = quantized.contiguous() + if not scales: + output = torch.empty_like(current, dtype=dtype) + _cast_kernel[grid]( + current, + output, + element_count, + block_size=_ELEMENT_BLOCK_SIZE, + ) + return output + + current_dtype = scales[0].dtype + for index, scale in enumerate(scales): + broadcast_scale = torch.broadcast_to(scale, shape) + if index: + current_dtype = torch.promote_types(current_dtype, scale.dtype) + output_dtype = dtype if index == len(scales) - 1 else current_dtype + output = torch.empty_like(current, dtype=output_dtype) + _multiply_broadcast_kernel[grid]( + current, + broadcast_scale, + output, + element_count, + shape=shape, + scale_strides=broadcast_scale.stride(), + block_size=_ELEMENT_BLOCK_SIZE, + ) + current = output + return current diff --git a/flashdreams/tests/accelerated/common/test_non_persistent_linear.py b/flashdreams/tests/accelerated/common/test_non_persistent_linear.py new file mode 100644 index 000000000..5eb53e8bc --- /dev/null +++ b/flashdreams/tests/accelerated/common/test_non_persistent_linear.py @@ -0,0 +1,42 @@ +# 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. + +"""Tests for shared accelerated inference modules.""" + +import pytest +import torch +import torch.nn.functional as F + +from flashdreams.accelerated.common.non_persistent_linear import ( + NonPersistentLinear, +) + +pytestmark = pytest.mark.ci_cpu + + +@pytest.mark.parametrize( + "bias", + (None, torch.tensor([0.25, -0.5])), + ids=("without-bias", "with-bias"), +) +def test_non_persistent_linear(bias: torch.Tensor | None) -> None: + """Apply a linear transformation without registering checkpoint parameters.""" + weight = torch.tensor(((1.0, 2.0, 3.0), (-1.0, 0.5, 2.0))) + inputs = torch.tensor(((1.0, 0.0, -1.0), (0.5, 2.0, 1.0))) + linear = NonPersistentLinear(weight, bias) + + torch.testing.assert_close(linear(inputs), F.linear(inputs, weight, bias)) + assert isinstance(linear, torch.nn.Linear) + assert list(linear.parameters()) == [] diff --git a/flashdreams/tests/accelerated/conftest.py b/flashdreams/tests/accelerated/conftest.py new file mode 100644 index 000000000..70e08be88 --- /dev/null +++ b/flashdreams/tests/accelerated/conftest.py @@ -0,0 +1,48 @@ +# 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 cuda_device() -> torch.device: + """Provide the active CUDA device. + + Returns: + Active CUDA device. + """ + if not torch.cuda.is_available(): + pytest.skip("CUDA required.") + return torch.device("cuda") + + +@pytest.fixture(scope="module") +def tma_device(cuda_device: torch.device) -> torch.device: + """Provide a CUDA device capable of launching TMA kernels. + + Args: + cuda_device: Active CUDA device. + + Returns: + Active CUDA device with compute capability 9.0 or newer. + """ + if torch.cuda.get_device_capability(cuda_device)[0] < 9: + pytest.skip("TMA kernels require compute capability 9.0 or newer.") + return cuda_device diff --git a/flashdreams/tests/accelerated/multi_head_attention/test_mha_optimized.py b/flashdreams/tests/accelerated/multi_head_attention/test_mha_optimized.py new file mode 100644 index 000000000..045e03c60 --- /dev/null +++ b/flashdreams/tests/accelerated/multi_head_attention/test_mha_optimized.py @@ -0,0 +1,571 @@ +# 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 parity tests for optimized and Torch multi-head attention.""" + +from __future__ import annotations + + +import pytest +import torch +from torch import Tensor + +from flashdreams.accelerated.multi_head_attention import ( + AttentionConfig, + AttentionType, + QKNormScope, + RoPEConfig, + RoPEScope, + RoPEStyle, +) +from flashdreams.accelerated.multi_head_attention.torch import TorchMultiHeadAttention +from flashdreams.accelerated.multi_head_attention.optimized import ( + QKVFusionOption, + QuantizationOption, + SDPABackend, + OptimizedImplConfig, + OptimizedHultiHeadAttention, +) +from flashdreams.accelerated.quantization.linear import QuantizedNonPersistentLinear +from flashdreams.accelerated.quantization.quantizer import ( + DTYPE_MAX, +) +from flashdreams.core.attention import BlockKVCache + +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 + +_ROPE_CASES = ( + (None, "none"), + *( + ( + RoPEConfig(style=style, scope=scope), + f"{style.value}-{scope.value}", + ) + for style in RoPEStyle + for scope in RoPEScope + ), +) +"""Supported rotary policies and their pytest identifiers.""" + +_ATTENTION_CONFIGS = tuple( + pytest.param( + AttentionConfig( + query_dim=_QUERY_DIM, + n_heads=_N_HEADS, + head_dim=_HEAD_DIM, + qk_norm_scope=qk_norm_scope, + rope_config=rope_config, + ), + id=f"norm-{qk_norm_scope.value}-rope-{rope_id}", + ) + for qk_norm_scope in QKNormScope + for rope_config, rope_id in _ROPE_CASES +) +"""Every supported normalization and rotary policy combination.""" + +_OPTIMIZED_IMPL_CONFIGS = tuple( + pytest.param( + OptimizedImplConfig( + qkv_fusion_option=qkv_fusion_option, + sdpa_backend=sdpa_backend, + use_tma=use_tma, + ), + id=f"{sdpa_backend.value}-{qkv_fusion_option.value}-{'tma' if use_tma else 'no-tma'}", + ) + for sdpa_backend in SDPABackend + for qkv_fusion_option in QKVFusionOption + for use_tma in (False, True) +) +"""Every supported optimized backend, fusion, and TMA preference combination.""" + + +class _AttentionModules: + """Provide checkpoint-compatible projection and normalization modules.""" + + attention_config: AttentionConfig + """Attention geometry and policies supplied by the concrete backend.""" + + @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 _initialize_modules(self) -> None: + """Initialize the shared checkpoint-compatible modules.""" + assert self.attention_config.context_dim is not None + self.q_proj = torch.nn.Linear( + self.attention_config.query_dim, self.attention_config.inner_dim, bias=True + ) + self.k_proj = torch.nn.Linear( + self.attention_config.context_dim, + self.attention_config.inner_dim, + bias=True, + ) + self.v_proj = torch.nn.Linear( + self.attention_config.context_dim, + self.attention_config.inner_dim, + bias=True, + ) + self.output_proj = torch.nn.Linear( + self.attention_config.inner_dim, self.attention_config.query_dim, bias=True + ) + if self.attention_config.qk_norm_scope is QKNormScope.NONE: + self.q_norm = torch.nn.Identity() + self.k_norm = torch.nn.Identity() + return + norm_dim = ( + self.attention_config.head_dim + if self.attention_config.qk_norm_scope is QKNormScope.HEAD + else self.attention_config.inner_dim + ) + self.q_norm = torch.nn.RMSNorm(norm_dim, eps=self.attention_config.qk_norm_eps) + self.k_norm = torch.nn.RMSNorm(norm_dim, eps=self.attention_config.qk_norm_eps) + + +class _TorchMHA(_AttentionModules, TorchMultiHeadAttention): + """Checkpoint-compatible Torch reference attention.""" + + def __init__( + self, + attention_type: AttentionType, + attention_config: AttentionConfig, + ) -> None: + """Initialize the Torch reference. + + Args: + attention_type: Relationship between query and context tokens. + attention_config: Shared attention geometry and policies. + """ + super().__init__(attention_type, attention_config) + self._initialize_modules() + + +class _OptimizedMHA(_AttentionModules, OptimizedHultiHeadAttention): + """Checkpoint-compatible Optimized attention under test.""" + + def __init__( + self, + attention_type: AttentionType, + attention_config: AttentionConfig, + optimized_impl_config: OptimizedImplConfig, + ) -> None: + """Initialize the Optimized implementation. + + Args: + attention_type: Relationship between query and context tokens. + attention_config: Shared attention geometry and policies. + optimized_impl_config: optimized backend and projection-fusion policies. + """ + super().__init__(attention_type, attention_config, optimized_impl_config) + self._initialize_modules() + self._initialize_derived_weights() + + +def _rope_freqs( + length: int, + attention_config: AttentionConfig, + generator: torch.Generator, + device: torch.device, +) -> Tensor | None: + """Generate rotary frequencies when the attention policy enables RoPE. + + Args: + length: Token sequence length. + attention_config: Shared attention policy. + generator: Seeded CUDA random generator. + device: CUDA device on which to allocate the frequencies. + + Returns: + Rotation angles shaped ``[L, 1, 1, D]``, or ``None`` when disabled. + """ + if attention_config.rope_config is None: + return None + half_freqs = torch.randn( + length, + 1, + 1, + attention_config.head_dim // 2, + generator=generator, + device=device, + dtype=torch.float32, + ) + if attention_config.rope_config.style is RoPEStyle.INTERLEAVED: + return half_freqs.repeat_interleave(2, dim=-1) + return torch.cat((half_freqs, half_freqs), dim=-1) + + +def _assert_close( + actual: Tensor, + expected: Tensor, + tolerance: float = 2e-2, +) -> None: + """Compare optimized output with the Torch reference.""" + torch.testing.assert_close(actual, expected, atol=tolerance, rtol=tolerance) + + +def _assert_cache_close( + actual: BlockKVCache, + expected: BlockKVCache, + tolerance: float = 2e-2, +) -> None: + """Compare visible optimized and Torch cache contents.""" + _assert_close( + actual.cached_k().to(expected.cached_k().dtype), + expected.cached_k(), + tolerance, + ) + _assert_close( + actual.cached_v().to(expected.cached_v().dtype), + expected.cached_v(), + tolerance, + ) + + +def _check_self_attention( + reference: _TorchMHA, + actual: _OptimizedMHA, + attention_config: AttentionConfig, + generator: torch.Generator, + device: torch.device, + tolerance: float = 2e-2, +) -> None: + """Compare streaming self-attention through cache fill and rolling. + + Args: + reference: Torch attention with weights shared by ``actual``. + actual: Optimized attention under test. + attention_config: Shared attention geometry and policies. + generator: Seeded CUDA random generator. + device: CUDA device on which to run the comparison. + tolerance: Absolute and relative comparison tolerance. + """ + reference_cache = reference.allocate_kv_cache( + batch_size=1, + chunk_size=_CHUNK_SIZE, + window_size=_WINDOW_SIZE, + sink_size=_SINK_SIZE, + device=device, + dtype=torch.bfloat16, + ) + actual_cache = actual.allocate_kv_cache( + batch_size=1, + chunk_size=_CHUNK_SIZE, + window_size=_WINDOW_SIZE, + sink_size=_SINK_SIZE, + device=device, + dtype=torch.bfloat16, + ) + expected_cache_dtype = ( + torch.float8_e4m3fn + if actual.optimized_impl_config.quantization.quantized_sdpa + else torch.bfloat16 + ) + assert actual_cache._k.dtype is expected_cache_dtype + assert actual_cache._v.dtype is expected_cache_dtype + for chunk_idx in range(3): + x = torch.randn( + 1, + _CHUNK_SIZE, + attention_config.query_dim, + generator=generator, + device=device, + dtype=torch.bfloat16, + ) + rope_length = ( + _SINK_SIZE + _WINDOW_SIZE + if attention_config.rope_config is not None + and attention_config.rope_config.scope is RoPEScope.AFTER_KV_CACHE + else _CHUNK_SIZE + ) + rope_freqs = _rope_freqs(rope_length, attention_config, generator, device) + + reference_cache.before_update(chunk_idx) + expected = reference(x, reference_cache, rope_freqs) + + actual_cache.before_update(chunk_idx) + output = actual(x, actual_cache, rope_freqs) + + _assert_close(output, expected, tolerance) + _assert_cache_close(actual_cache, reference_cache, tolerance) + reference_cache.after_update(chunk_idx) + actual_cache.after_update(chunk_idx) + + +def _check_cross_attention( + reference: _TorchMHA, + actual: _OptimizedMHA, + attention_config: AttentionConfig, + generator: torch.Generator, + device: torch.device, + tolerance: float = 2e-2, +) -> None: + """Compare static cross-attention with the Torch reference. + + Args: + reference: Torch attention with weights shared by ``actual``. + actual: Optimized attention under test. + attention_config: Shared attention geometry and policies. + generator: Seeded CUDA random generator. + device: CUDA device on which to run the comparison. + tolerance: Absolute and relative comparison tolerance. + """ + assert attention_config.context_dim is not None + context = torch.randn( + 1, + 24, + attention_config.context_dim, + generator=generator, + device=device, + dtype=torch.bfloat16, + ) + query_length = ( + context.shape[-2] + if attention_config.rope_config is not None + and attention_config.rope_config.scope is RoPEScope.AFTER_KV_CACHE + else 8 + ) + query = torch.randn( + 1, + query_length, + attention_config.query_dim, + generator=generator, + device=device, + dtype=torch.bfloat16, + ) + context_rope = _rope_freqs(24, attention_config, generator, device) + query_rope = _rope_freqs(query_length, attention_config, generator, device) + reference_cache = reference.compute_kv(context, context_rope) + actual_cache = actual.compute_kv(context, context_rope) + expected_cache_dtype = ( + torch.float8_e4m3fn + if actual.optimized_impl_config.quantization.quantized_sdpa + else torch.bfloat16 + ) + assert actual_cache._k.dtype is expected_cache_dtype + assert actual_cache._v.dtype is expected_cache_dtype + + expected = reference(query, reference_cache, query_rope) + output = actual(query, actual_cache, query_rope) + + _assert_close(output, expected, tolerance) + _assert_cache_close(actual_cache, reference_cache, tolerance) + + +@pytest.mark.parametrize( + "attention_type", tuple(AttentionType), ids=lambda value: value.value +) +@pytest.mark.parametrize("attention_config", _ATTENTION_CONFIGS) +@pytest.mark.parametrize("optimized_impl_config", _OPTIMIZED_IMPL_CONFIGS) +@torch.inference_mode() +def test_mha_optimized_matches_torch( + cuda_device: torch.device, + attention_type: AttentionType, + attention_config: AttentionConfig, + optimized_impl_config: OptimizedImplConfig, +) -> None: + """Match Optimized attention with Torch for every supported policy.""" + torch.manual_seed(7) + reference = _TorchMHA(attention_type, attention_config) + actual = _OptimizedMHA(attention_type, attention_config, optimized_impl_config) + actual.load_state_dict(reference.state_dict(), strict=True) + reference.to(device=cuda_device, dtype=torch.bfloat16).eval() + actual.to(device=cuda_device, dtype=torch.bfloat16).eval() + + generator = torch.Generator(device=cuda_device).manual_seed(11) + if attention_type is AttentionType.SELF_ATTENTION: + _check_self_attention( + reference, actual, attention_config, generator, cuda_device + ) + else: + _check_cross_attention( + reference, actual, attention_config, generator, cuda_device + ) + + +@pytest.mark.parametrize( + "attention_type", tuple(AttentionType), ids=lambda value: value.value +) +@pytest.mark.parametrize( + "qkv_fusion_option", tuple(QKVFusionOption), ids=lambda value: value.value +) +@pytest.mark.parametrize( + "projection_dtype", + tuple(DTYPE_MAX), + ids=lambda dtype: str(dtype).removeprefix("torch."), +) +@torch.inference_mode() +def test_mha_optimized_quantized_projections_match_torch( + cuda_device: torch.device, + attention_type: AttentionType, + qkv_fusion_option: QKVFusionOption, + projection_dtype: torch.dtype, +) -> None: + """Match quantized Q/K/V projections against native-precision attention.""" + attention_config = AttentionConfig( + query_dim=_QUERY_DIM, + n_heads=_N_HEADS, + head_dim=_HEAD_DIM, + qk_norm_scope=QKNormScope.HEAD, + ) + optimized_impl_config = OptimizedImplConfig( + qkv_fusion_option=qkv_fusion_option, + quantization=QuantizationOption(projection=projection_dtype), + sdpa_backend=SDPABackend.CUDNN, + use_tma=False, + ) + torch.manual_seed(17) + reference = _TorchMHA(attention_type, attention_config) + actual = _OptimizedMHA(attention_type, attention_config, optimized_impl_config) + actual.load_state_dict(reference.state_dict(), strict=True) + reference.to(device=cuda_device, dtype=torch.bfloat16).eval() + actual.to(device=cuda_device, dtype=torch.bfloat16).eval() + + assert set(actual.state_dict()) == set(reference.state_dict()) + assert isinstance(actual.quantized_query_projection, QuantizedNonPersistentLinear) + assert actual.quantized_query_projection.dtype is projection_dtype + if qkv_fusion_option is QKVFusionOption.NONE: + assert isinstance(actual.quantized_key_projection, QuantizedNonPersistentLinear) + assert isinstance( + actual.quantized_value_projection, QuantizedNonPersistentLinear + ) + assert actual.fused_qkv is None + assert actual.fused_kv is None + else: + assert actual.quantized_key_projection is None + assert actual.quantized_value_projection is None + assert isinstance(actual.fused_kv, QuantizedNonPersistentLinear) + assert actual.fused_kv.dtype is projection_dtype + if qkv_fusion_option is QKVFusionOption.FULL: + assert isinstance(actual.fused_qkv, QuantizedNonPersistentLinear) + assert actual.fused_qkv.dtype is projection_dtype + else: + assert actual.fused_qkv is None + + quantization_tolerance = ( + torch.finfo(projection_dtype).eps + if projection_dtype.is_floating_point + else 5 / DTYPE_MAX[projection_dtype] + ) + tolerance = max(2e-2, quantization_tolerance) + generator = torch.Generator(device=cuda_device).manual_seed(19) + if attention_type is AttentionType.SELF_ATTENTION: + _check_self_attention( + reference, + actual, + attention_config, + generator, + cuda_device, + tolerance, + ) + else: + _check_cross_attention( + reference, + actual, + attention_config, + generator, + cuda_device, + tolerance, + ) + + +@pytest.mark.parametrize( + "attention_type", tuple(AttentionType), ids=lambda value: value.value +) +@pytest.mark.parametrize("rope_scope", tuple(RoPEScope), ids=lambda value: value.value) +@pytest.mark.parametrize( + "sdpa_backend", tuple(SDPABackend), ids=lambda value: value.value +) +@pytest.mark.parametrize("use_tma", (False, True), ids=("no-tma", "tma")) +@torch.inference_mode() +def test_mha_optimized_quantized_sdpa_matches_torch( + cuda_device: torch.device, + attention_type: AttentionType, + rope_scope: RoPEScope, + sdpa_backend: SDPABackend, + use_tma: bool, +) -> None: + """Exercise the unscaled e4m3 SDPA/cache contract across backends.""" + attention_config = AttentionConfig( + query_dim=_QUERY_DIM, + n_heads=_N_HEADS, + head_dim=_HEAD_DIM, + qk_norm_scope=QKNormScope.HEAD, + rope_config=RoPEConfig(style=RoPEStyle.SPLIT, scope=rope_scope), + ) + optimized_impl_config = OptimizedImplConfig( + qkv_fusion_option=QKVFusionOption.FULL, + quantization=QuantizationOption(quantized_sdpa=True), + sdpa_backend=sdpa_backend, + use_tma=use_tma, + ) + torch.manual_seed(23) + reference = _TorchMHA(attention_type, attention_config) + actual = _OptimizedMHA(attention_type, attention_config, optimized_impl_config) + actual.load_state_dict(reference.state_dict(), strict=True) + reference.to(device=cuda_device, dtype=torch.bfloat16).eval() + actual.to(device=cuda_device, dtype=torch.bfloat16).eval() + + generator = torch.Generator(device=cuda_device).manual_seed(29) + tolerance = 2 * torch.finfo(torch.float8_e4m3fn).eps + if attention_type is AttentionType.SELF_ATTENTION: + _check_self_attention( + reference, + actual, + attention_config, + generator, + cuda_device, + tolerance, + ) + else: + _check_cross_attention( + reference, + actual, + attention_config, + generator, + cuda_device, + tolerance, + ) diff --git a/flashdreams/tests/accelerated/multi_head_attention/test_mha_torch.py b/flashdreams/tests/accelerated/multi_head_attention/test_mha_torch.py new file mode 100644 index 000000000..360dec73e --- /dev/null +++ b/flashdreams/tests/accelerated/multi_head_attention/test_mha_torch.py @@ -0,0 +1,172 @@ +# 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. + +"""Correctness tests for Torch multi-head attention.""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn.functional as F +from torch import Tensor, nn + +from flashdreams.accelerated.multi_head_attention import ( + AttentionConfig, + AttentionType, + QKNormScope, + RoPEConfig, + RoPEScope, + RoPEStyle, +) +from flashdreams.accelerated.multi_head_attention.torch import TorchMultiHeadAttention + +pytestmark = pytest.mark.ci_cpu + + +class _IdentityMHA(TorchMultiHeadAttention): + """Provide identity projections for direct attention comparisons.""" + + @property + def query_projection(self) -> nn.Linear: + """Return the query projection.""" + return self.projection + + @property + def key_projection(self) -> nn.Linear: + """Return the key projection.""" + return self.projection + + @property + def value_projection(self) -> nn.Linear: + """Return the value projection.""" + return self.projection + + @property + def output_projection(self) -> nn.Linear: + """Return the output projection.""" + return self.projection + + @property + def query_norm(self) -> nn.Module: + """Return identity query normalization.""" + return self.norm + + @property + def key_norm(self) -> nn.Module: + """Return identity key normalization.""" + return self.norm + + def __init__(self, rope_style: RoPEStyle) -> None: + """Initialize identity attention with after-cache RoPE. + + Args: + rope_style: Feature pairing convention under test. + """ + super().__init__( + AttentionType.SELF_ATTENTION, + AttentionConfig( + query_dim=4, + n_heads=1, + head_dim=4, + qk_norm_scope=QKNormScope.NONE, + rope_config=RoPEConfig( + style=rope_style, + scope=RoPEScope.AFTER_KV_CACHE, + ), + ), + ) + self.projection = nn.Linear(4, 4, bias=False) + self.norm = nn.Identity() + with torch.no_grad(): + self.projection.weight.copy_(torch.eye(4)) + + +def _apply_rope(x: Tensor, rope_freqs: Tensor, style: RoPEStyle) -> Tensor: + """Apply RoPE independently for the expected attention result. + + Args: + x: Projected query or key heads. + rope_freqs: Rotation angles for every token in ``x``. + style: Feature pairing convention under test. + + Returns: + Rotated projected heads. + """ + freqs = rope_freqs[:, 0, 0].reshape(1, x.shape[1], 1, x.shape[-1]) + if style is RoPEStyle.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() + rotated * freqs.sin() + + +@pytest.mark.parametrize("rope_style", tuple(RoPEStyle), ids=lambda value: value.value) +@torch.inference_mode() +def test_after_kv_cache_rope_matches_visible_cache_positions( + rope_style: RoPEStyle, +) -> None: + """Rotate visible keys after cache fill and rolling without changing storage. + + Args: + rope_style: Feature pairing convention under test. + """ + attention = _IdentityMHA(rope_style) + cache = attention.allocate_kv_cache( + batch_size=1, + chunk_size=2, + window_size=4, + sink_size=0, + device="cpu", + dtype=torch.float32, + ) + chunks = ( + torch.tensor([[[1.0, 2.0, 3.0, 4.0], [2.0, -1.0, 0.5, 3.0]]]), + torch.tensor([[[0.5, 1.5, -2.0, 1.0], [3.0, 0.5, 2.0, -1.0]]]), + torch.tensor([[[-1.0, 2.5, 1.0, 0.5], [1.5, -0.5, 3.0, 2.0]]]), + ) + rope_freqs = torch.linspace(0.1, 1.6, steps=16).reshape(4, 1, 1, 4) + + for chunk_idx, chunk in enumerate(chunks): + cache.before_update(chunk_idx) + write_end = cache.write_end + write_start = write_end - cache.chunk_size + actual = attention(chunk, cache, rope_freqs) + + visible = torch.cat(chunks[max(0, chunk_idx - 1) : chunk_idx + 1], dim=1) + query = _apply_rope( + chunk.reshape(1, 2, 1, 4), + rope_freqs[write_start:write_end], + rope_style, + ) + key = _apply_rope( + visible.reshape(1, visible.shape[1], 1, 4), + rope_freqs[: visible.shape[1]], + rope_style, + ) + value = visible.reshape(1, visible.shape[1], 1, 4) + expected = ( + F.scaled_dot_product_attention( + query.transpose(1, 2), + key.transpose(1, 2), + value.transpose(1, 2), + ) + .transpose(1, 2) + .flatten(-2) + ) + + torch.testing.assert_close(actual, expected) + torch.testing.assert_close(cache.cached_k(), value) + cache.after_update(chunk_idx) diff --git a/flashdreams/tests/accelerated/multi_head_attention/triton/test_flash_attention_2_kernel.py b/flashdreams/tests/accelerated/multi_head_attention/triton/test_flash_attention_2_kernel.py new file mode 100644 index 000000000..e8fa3179b --- /dev/null +++ b/flashdreams/tests/accelerated/multi_head_attention/triton/test_flash_attention_2_kernel.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 FlashAttention2 kernels.""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn.functional as F + +from flashdreams.accelerated.multi_head_attention.triton import flash_attention_2 + +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_flash_attention_matches_sdpa( + cuda_device: torch.device, + query_length: int, + key_length: int, + head_dim: int, +) -> None: + """Match pointer-based 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: + cuda_device: Active CUDA device. + 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=cuda_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=cuda_device, + dtype=torch.bfloat16, + ) + key = torch.randn( + 1, + key_length, + 2, + head_dim, + generator=generator, + device=cuda_device, + dtype=torch.bfloat16, + ) + value = torch.randn( + key.shape, + generator=generator, + device=cuda_device, + dtype=torch.bfloat16, + ) + + actual = flash_attention_2(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/multi_head_attention/triton/test_flash_attention_2_tma_kernel.py b/flashdreams/tests/accelerated/multi_head_attention/triton/test_flash_attention_2_tma_kernel.py new file mode 100644 index 000000000..e072a6565 --- /dev/null +++ b/flashdreams/tests/accelerated/multi_head_attention/triton/test_flash_attention_2_tma_kernel.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.multi_head_attention.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/quantization/test_linear.py b/flashdreams/tests/accelerated/quantization/test_linear.py new file mode 100644 index 000000000..3c1147f3e --- /dev/null +++ b/flashdreams/tests/accelerated/quantization/test_linear.py @@ -0,0 +1,150 @@ +# 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 quantized nonpersistent linear transformations.""" + +import pytest +import torch +import torch.nn.functional as F + +from flashdreams.accelerated.common.non_persistent_linear import ( + NonPersistentLinear, +) +from flashdreams.accelerated.quantization.linear import ( + QuantizedNonPersistentLinear, + WeightGranularity, +) +from flashdreams.accelerated.quantization.quantizer import ( + DTYPE_MAX, + Granularity, + quantize, +) + +pytestmark = pytest.mark.ci_cpu + + +@pytest.mark.parametrize("dtype", DTYPE_MAX) +@pytest.mark.parametrize("weight_granularity", WeightGranularity) +@pytest.mark.parametrize("input_granularity", Granularity) +@pytest.mark.parametrize("use_bias", (False, True), ids=("without-bias", "with-bias")) +def test_quantized_non_persistent_linear( + dtype: torch.dtype, + weight_granularity: WeightGranularity, + input_granularity: Granularity, + use_bias: bool, +) -> None: + """Match quantized projection paths against full-precision linear inference.""" + generator = torch.Generator().manual_seed(4) + weight = torch.randn((16, 32), generator=generator) + bias = torch.randn((16,), generator=generator) if use_bias else None + inputs = torch.randn((2, 3, 32), generator=generator) + linear = QuantizedNonPersistentLinear(weight, bias, weight_granularity, dtype) + + quantized, scale = quantize(inputs, dtype, input_granularity, axis=-1) + dynamic_output = linear(inputs, input_granularity) + prequantized_output = linear(quantized, scale) + + assert dynamic_output.shape == (2, 3, 16) + assert dynamic_output.dtype is torch.float16 + torch.testing.assert_close(dynamic_output, prequantized_output, rtol=0, atol=0) + + full_precision_output = F.linear(inputs, weight, bias) + tolerance = ( + torch.finfo(dtype).eps if dtype.is_floating_point else 2 / DTYPE_MAX[dtype] + ) + torch.testing.assert_close( + dynamic_output.float(), + full_precision_output, + rtol=tolerance, + atol=tolerance * full_precision_output.abs().amax().item(), + ) + relative_error = ( + dynamic_output.float() - full_precision_output + ).norm() / full_precision_output.norm() + assert relative_error.item() < tolerance + + expected_weight_dtype = torch.float8_e4m3fn if dtype is torch.float8_e5m2 else dtype + assert linear.dtype is dtype + assert linear.weight.dtype is expected_weight_dtype + + +@pytest.mark.parametrize("out_dtype", (torch.bfloat16, torch.float32)) +def test_quantized_linear_output_dtype(out_dtype: torch.dtype) -> None: + """Return projected activations in the requested data type.""" + linear = QuantizedNonPersistentLinear( + torch.eye(16), + torch.ones(16), + WeightGranularity.PER_OUT_CHANNEL, + torch.float8_e4m3fn, + ) + + output = linear(torch.ones((2, 16)), Granularity.SLICE, out_dtype=out_dtype) + + assert output.dtype is out_dtype + torch.testing.assert_close(output, torch.full_like(output, 2)) + + +def test_quantized_linear_buffers_are_nonpersistent() -> None: + """Keep derived quantized tensors out of parameters and checkpoints.""" + linear = QuantizedNonPersistentLinear( + torch.eye(16), + torch.ones(16), + WeightGranularity.TENSOR, + torch.int8, + ) + + assert isinstance(linear, NonPersistentLinear) + assert list(linear.parameters()) == [] + assert linear.state_dict() == {} + assert set(dict(linear.named_buffers())) == {"weight", "bias", "weight_scale"} + + +@pytest.mark.parametrize("granularity", Granularity) +def test_quantized_linear_empty_input(granularity: Granularity) -> None: + """Preserve empty leading dimensions without reducing an empty tensor.""" + linear = QuantizedNonPersistentLinear( + torch.eye(16), + None, + WeightGranularity.TENSOR, + torch.float8_e4m3fn, + ) + + output = linear(torch.empty((0, 3, 16)), granularity) + + assert output.shape == (0, 3, 16) + assert output.dtype is torch.float16 + + +def test_quantized_linear_rejects_invalid_prequantized_input() -> None: + """Reject prequantized activations whose dtype or scale is incompatible.""" + linear = QuantizedNonPersistentLinear( + torch.eye(16), + None, + WeightGranularity.TENSOR, + torch.float8_e4m3fn, + ) + inputs = torch.ones((2, 3, 16)) + quantized, scale = quantize(inputs, torch.float8_e4m3fn, Granularity.SLICE, axis=-1) + + with pytest.raises(ValueError, match="input dtype"): + linear(quantized.to(torch.float8_e5m2), scale) + with pytest.raises(ValueError, match="scale shape"): + linear(quantized, torch.ones((2, 1))) + with pytest.raises(ValueError, match="FP32 scale"): + linear(quantized, scale.to(torch.float16)) + with pytest.raises(ValueError, match="last dim"): + linear(quantized[..., :-1], scale) + with pytest.raises(ValueError, match="scale tensor or Granularity"): + linear(inputs, "slice") # type: ignore[call-overload] diff --git a/flashdreams/tests/accelerated/quantization/test_quantizer.py b/flashdreams/tests/accelerated/quantization/test_quantizer.py new file mode 100644 index 000000000..bf96e5316 --- /dev/null +++ b/flashdreams/tests/accelerated/quantization/test_quantizer.py @@ -0,0 +1,306 @@ +# 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. + +"""CUDA correctness tests for Torch and Triton tensor quantization.""" + +import pytest +import torch + +from flashdreams.accelerated.quantization.quantizer import ( + DTYPE_MAX, + Granularity, + dequantize, + quantize, +) + +pytestmark = pytest.mark.ci_gpu + +_IMPLEMENTATIONS = ( + pytest.param(False, id="torch"), + pytest.param(True, id="triton"), +) + +_GRANULARITY_AXES = ( + pytest.param(Granularity.TENSOR, -1, id="tensor"), + pytest.param(Granularity.SLICE, 0, id="slice-axis-0"), + pytest.param(Granularity.SLICE, 1, id="slice-axis-1"), + pytest.param(Granularity.SLICE, -1, id="slice-axis-negative-1"), +) + + +@pytest.mark.parametrize("use_triton", _IMPLEMENTATIONS) +def test_quantize_preserves_reduced_dimensions( + cuda_device: torch.device, use_triton: bool +) -> None: + original = torch.linspace(-4.0, 4.0, 8**4, device=cuda_device).reshape(8, 8, 8, 8) + + tensor_quantized, tensor_scale = quantize( + original, torch.float8_e4m3fn, Granularity.TENSOR, use_triton=use_triton + ) + slice_quantized, slice_scale = quantize( + original, + torch.float8_e4m3fn, + Granularity.SLICE, + axis=2, + use_triton=use_triton, + ) + + assert tensor_scale.shape == (1, 1, 1, 1) + assert slice_scale.shape == (8, 8, 1, 8) + assert tensor_quantized.dtype is torch.float8_e4m3fn + assert slice_quantized.dtype is torch.float8_e4m3fn + torch.testing.assert_close( + tensor_scale, + original.abs().amax().reshape(1, 1, 1, 1) / DTYPE_MAX[torch.float8_e4m3fn], + ) + torch.testing.assert_close( + slice_scale, + original.abs().amax(dim=2, keepdim=True) / DTYPE_MAX[torch.float8_e4m3fn], + ) + torch.testing.assert_close( + dequantize( + slice_quantized, + slice_scale, + dtype=torch.float32, + use_triton=use_triton, + ), + original, + rtol=0.06, + atol=0.01, + ) + + +@pytest.mark.parametrize("use_triton", _IMPLEMENTATIONS) +def test_zero_groups_and_multiple_dequantization_scales( + cuda_device: torch.device, use_triton: bool +) -> None: + zeros = torch.zeros(2, 3, device=cuda_device) + zero_quantized, zero_scale = quantize( + zeros, + torch.float8_e4m3fn, + Granularity.SLICE, + use_triton=use_triton, + ) + + assert zero_scale.shape == (2, 1) + assert torch.isfinite(zero_scale).all() + torch.testing.assert_close( + dequantize( + zero_quantized, + zero_scale, + dtype=torch.float32, + use_triton=use_triton, + ), + zeros, + ) + + quantized = torch.tensor( + [[1.0, -2.0], [3.0, -4.0]], + device=cuda_device, + dtype=torch.float8_e4m3fn, + ) + dequantized = dequantize( + quantized, + torch.tensor([[0.5], [2.0]], device=cuda_device), + torch.tensor([[2.0, 4.0]], device=cuda_device), + torch.tensor(4.0, device=cuda_device), + use_triton=use_triton, + ) + assert dequantized.dtype is torch.float16 + torch.testing.assert_close( + dequantized, + torch.tensor( + [[4.0, -16.0], [48.0, -128.0]], + device=cuda_device, + dtype=torch.float16, + ), + ) + torch.testing.assert_close( + dequantize( + quantized, + dtype=torch.float32, + use_triton=use_triton, + ), + quantized.float(), + ) + + +@pytest.mark.parametrize("use_triton", _IMPLEMENTATIONS) +def test_int8_quantization_rounds_to_nearest_integer( + cuda_device: torch.device, use_triton: bool +) -> None: + original = torch.tensor([[-2.0, -1.0, 0.0, 1.0, 2.0]], device=cuda_device) + + quantized, scale = quantize( + original, torch.int8, Granularity.SLICE, use_triton=use_triton + ) + + assert DTYPE_MAX[torch.int8] == torch.iinfo(torch.int8).max + assert quantized.dtype is torch.int8 + assert scale.shape == (1, 1) + torch.testing.assert_close( + quantized, + torch.tensor( + [[-127, -64, 0, 64, 127]], + device=cuda_device, + dtype=torch.int8, + ), + ) + torch.testing.assert_close( + dequantize( + quantized, + scale, + dtype=torch.float32, + use_triton=use_triton, + ), + original, + rtol=0.0, + atol=scale.item() / 2, + ) + + +@pytest.mark.parametrize("dtype", DTYPE_MAX) +@pytest.mark.parametrize("granularity,axis", _GRANULARITY_AXES) +@pytest.mark.parametrize("use_triton", _IMPLEMENTATIONS) +def test_random_tensor_quantize_dequantize_round_trip( + cuda_device: torch.device, + dtype: torch.dtype, + granularity: Granularity, + axis: int, + use_triton: bool, +) -> None: + generator = torch.Generator(device=cuda_device).manual_seed(0) + original = torch.randn((6, 5, 4), device=cuda_device, generator=generator).permute( + 2, 1, 0 + ) + + quantized, scale = quantize( + original, + dtype, + granularity, + axis=axis, + use_triton=use_triton, + ) + restored = dequantize( + quantized, + scale, + dtype=torch.float32, + use_triton=use_triton, + ) + + reduction_axis: int | tuple[int, ...] + reduction_axis = ( + tuple(range(original.ndim)) if granularity is Granularity.TENSOR else axis + ) + expected_scale = ( + original.float().abs().amax(dim=reduction_axis, keepdim=True) / DTYPE_MAX[dtype] + ).clamp_min(torch.finfo(torch.float32).tiny) + torch.testing.assert_close(scale, expected_scale) + + if dtype.is_floating_point: + rtol = torch.finfo(dtype).eps + atol = scale.max().item() + else: + rtol = 0.0 + atol = scale.max().item() / 2 + torch.testing.assert_close(restored, original, rtol=rtol, atol=atol) + + +@pytest.mark.parametrize("use_triton", _IMPLEMENTATIONS) +def test_quantize_reduction_larger_than_one_tile( + cuda_device: torch.device, use_triton: bool +) -> None: + generator = torch.Generator(device=cuda_device).manual_seed(0) + original = torch.randn( + (2, 20_001), + device=cuda_device, + dtype=torch.float16, + generator=generator, + ) + quantized, scale = quantize( + original, + torch.int8, + Granularity.SLICE, + axis=1, + use_triton=use_triton, + ) + expected_quantized, expected_scale = quantize( + original, torch.int8, Granularity.SLICE, axis=1, use_triton=False + ) + + torch.testing.assert_close(scale, expected_scale, rtol=0, atol=0) + torch.testing.assert_close(quantized, expected_quantized, rtol=0, atol=0) + + +@pytest.mark.parametrize("dtype", DTYPE_MAX) +@pytest.mark.parametrize("granularity", Granularity) +@pytest.mark.parametrize("use_triton", _IMPLEMENTATIONS) +def test_quantized_gemm( + cuda_device: torch.device, + dtype: torch.dtype, + granularity: Granularity, + use_triton: bool, +) -> None: + generator = torch.Generator(device=cuda_device).manual_seed(1) + left = torch.rand((32, 32), device=cuda_device, generator=generator) + right = torch.rand((32, 32), device=cuda_device, generator=generator) + + left_quantized, left_scale = quantize( + left, dtype, granularity, axis=1, use_triton=use_triton + ) + if dtype is torch.int8: + right_quantized, right_scale = quantize( + right, dtype, granularity, axis=0, use_triton=use_triton + ) + quantized_product = torch._int_mm(left_quantized, right_quantized) + restored = dequantize( + quantized_product, + left_scale, + right_scale, + dtype=torch.float32, + use_triton=use_triton, + ) + else: + right_dtype = torch.float8_e4m3fn if dtype is torch.float8_e5m2 else dtype + right_quantized_t, right_scale_t = quantize( + right.T.contiguous(), + right_dtype, + granularity, + axis=1, + use_triton=use_triton, + ) + right_quantized = right_quantized_t.T + right_scale = right_scale_t.T + scaled_out_dtype = ( + torch.bfloat16 if granularity is Granularity.SLICE else torch.float32 + ) + restored = torch._scaled_mm( + left_quantized, + right_quantized, + left_scale, + right_scale, + out_dtype=scaled_out_dtype, + ).float() + expected = left @ right + + if granularity is Granularity.SLICE: + assert left_scale.shape == (32, 1) + assert right_scale.shape == (1, 32) + relative_error = (restored - expected).norm() / expected.norm() + if dtype.is_floating_point: + tolerance = torch.finfo(dtype).eps + else: + tolerance = 1.0 / DTYPE_MAX[dtype] + assert relative_error.item() < tolerance diff --git a/integrations/omnidreams/README.md b/integrations/omnidreams/README.md index be0561bce..74d5bc7d3 100644 --- a/integrations/omnidreams/README.md +++ b/integrations/omnidreams/README.md @@ -178,6 +178,109 @@ 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. +## 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`, `optimized_cudnn`, and `optimized_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`, `optimized_cudnn`, `optimized_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`, `optimized_cudnn`, `optimized_fa2`, and native `cuda` + implementations at the runner's production 704x1280 resolution and scheduler + configuration. + +The selected optimized configurations use row-scaled FP8 projections. `optimized_cudnn` +uses PyTorch's cuDNN SDPA backend with a BF16 self-attention cache, while +`optimized_fa2` uses Triton FlashAttention2 (FA2) with an E4M3 cache. Text and +cross-view attention remain BF16. + +The optimized 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 (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..e2c7ebe1b --- /dev/null +++ b/integrations/omnidreams/benchmarks/cases.py @@ -0,0 +1,188 @@ +# 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.optimized import ( + QKVFusionOption, + QuantizationOption, + SDPABackend, + OptimizedImplConfig, +) + + +@dataclass(frozen=True) +class AttentionBenchmarkCase: + """Configuration for one attention benchmark implementation.""" + + implementation: str + """Stable implementation name used by pytest and pipeline setup.""" + + self_attention_backend: AttentionBackend + """Self-attention implementation configured for this case.""" + + cross_attention_backend: AttentionBackend + """Cross-attention implementation configured for this case.""" + + self_attn_optimized_impl_config: OptimizedImplConfig = OptimizedImplConfig( + qkv_fusion_option=QKVFusionOption.FULL, + sdpa_backend=SDPABackend.FA2, + ) + """Optimized implementation policy used by accelerated self-attention.""" + + cross_attn_optimized_impl_config: OptimizedImplConfig = OptimizedImplConfig( + qkv_fusion_option=QKVFusionOption.FUSE_KV, + sdpa_backend=SDPABackend.FA2, + ) + """Optimized implementation 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 = [ + # Pytorch reference implementation. + AttentionBenchmarkCase( + implementation="omnidreams_torch", + self_attention_backend=AttentionBackend.OMNIDREAMS, + cross_attention_backend=AttentionBackend.OMNIDREAMS, + self_attn_optimized_impl_config=OptimizedImplConfig( + qkv_fusion_option=QKVFusionOption.NONE, + sdpa_backend=SDPABackend.CUDNN, + ), + cross_attn_optimized_impl_config=OptimizedImplConfig( + qkv_fusion_option=QKVFusionOption.NONE, + sdpa_backend=SDPABackend.CUDNN, + ), + ), + # Selected production-shaped pair from the recorded GB300 MHA benchmark. + AttentionBenchmarkCase( + implementation="optimized_cudnn_fp8_self_full_no_tma_cross_none_tma", + self_attention_backend=AttentionBackend.OPTIMIZED, + cross_attention_backend=AttentionBackend.OPTIMIZED, + self_attn_optimized_impl_config=OptimizedImplConfig( + qkv_fusion_option=QKVFusionOption.FULL, + sdpa_backend=SDPABackend.CUDNN, + use_tma=False, + quantization=QuantizationOption(projection=torch.float8_e4m3fn), + ), + cross_attn_optimized_impl_config=OptimizedImplConfig( + qkv_fusion_option=QKVFusionOption.NONE, + sdpa_backend=SDPABackend.CUDNN, + use_tma=True, + quantization=QuantizationOption(projection=torch.float8_e4m3fn), + ), + ), + # RTX PRO 6000 quantized SDPA pair with e4m3 projections enabled. + AttentionBenchmarkCase( + implementation="optimized_fa2_quantized_sdpa_self_full_tma_cross_none_tma", + self_attention_backend=AttentionBackend.OPTIMIZED, + cross_attention_backend=AttentionBackend.OPTIMIZED, + self_attn_optimized_impl_config=OptimizedImplConfig( + qkv_fusion_option=QKVFusionOption.FULL, + sdpa_backend=SDPABackend.FA2, + use_tma=True, + quantization=QuantizationOption( + projection=torch.float8_e4m3fn, + quantized_sdpa=True, + ), + ), + cross_attn_optimized_impl_config=OptimizedImplConfig( + qkv_fusion_option=QKVFusionOption.NONE, + sdpa_backend=SDPABackend.FA2, + use_tma=True, + quantization=QuantizationOption( + projection=torch.float8_e4m3fn, + quantized_sdpa=True, + ), + ), + minimum_compute_capability=(9, 0), + ), + # Native CUDA implementation. + AttentionBenchmarkCase( + implementation="cuda", + self_attention_backend=AttentionBackend.OMNIDREAMS, + cross_attention_backend=AttentionBackend.OMNIDREAMS, + native_dit=True, + ), + AttentionBenchmarkCase( + implementation="cuda_sparge", + self_attention_backend=AttentionBackend.OMNIDREAMS, + cross_attention_backend=AttentionBackend.OMNIDREAMS, + 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, + 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, + 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.implementation} 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..2c94ea795 --- /dev/null +++ b/integrations/omnidreams/benchmarks/test_modules.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. + +"""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.optimized import ( + QKVFusionOption, + QuantizationOption, + SDPABackend, + OptimizedImplConfig, +) +from flashdreams.core.attention.rope import RotaryPositionEmbedding3D +from integrations.omnidreams.benchmarks.cases import ( + 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 + + +def _implementation_id(optimized_impl_config: OptimizedImplConfig | None) -> str: + """Return a stable pytest identifier for an attention implementation.""" + if optimized_impl_config is None: + return "omnidreams" + backend = optimized_impl_config.sdpa_backend.value + fusion = optimized_impl_config.qkv_fusion_option.value.replace("_", "-") + tma = "tma" if optimized_impl_config.use_tma else "no-tma" + projection_dtype = optimized_impl_config.quantization.projection + quantization = ( + "" + if projection_dtype is None + else f"-projection-{projection_dtype}".replace("torch.", "").replace("_", "-") + ) + quantized_sdpa = ( + "-quantized-sdpa" + if optimized_impl_config.quantization.quantized_sdpa + else "" + ) + return f"optimized-{backend}-{fusion}-{tma}{quantization}{quantized_sdpa}" + + +_OPTIMIZED_IMPL_CONFIGS = ( + *( + OptimizedImplConfig( + qkv_fusion_option=qkv_fusion_option, + quantization=QuantizationOption(projection=projection_dtype), + sdpa_backend=sdpa_backend, + use_tma=use_tma, + ) + for sdpa_backend in SDPABackend + for qkv_fusion_option in QKVFusionOption + for use_tma in (False, True) + for projection_dtype in (None, torch.float8_e4m3fn) + ), + *( + OptimizedImplConfig( + qkv_fusion_option=qkv_fusion_option, + quantization=QuantizationOption( + projection=torch.float8_e4m3fn, + quantized_sdpa=True, + ), + sdpa_backend=sdpa_backend, + use_tma=use_tma, + ) + for sdpa_backend in SDPABackend + for qkv_fusion_option in QKVFusionOption + for use_tma in (False, True) + ), +) +"""Every optimized SDPA, fusion, TMA, and attention-quantization policy.""" + +_MODULE_SELF_ATTENTION_CONFIGS = (None, *_OPTIMIZED_IMPL_CONFIGS) +"""Reference self-attention plus every Optimized implementation config.""" + +_MODULE_CROSS_ATTENTION_CONFIGS = ( + None, + *( + config + for config in _OPTIMIZED_IMPL_CONFIGS + if config.qkv_fusion_option is not QKVFusionOption.FULL + ), +) +"""Reference cross-attention plus every valid Optimized implementation config. + +Full QKV fusion requires equal query and context widths, which production +OmniDreams text cross-attention does not have. +""" + + +def _block_case( + self_config: OptimizedImplConfig | None, + cross_config: OptimizedImplConfig | None, +) -> AttentionBenchmarkCase: + """Build one self/cross implementation combination for the DiT block.""" + reference_config = OptimizedImplConfig( + qkv_fusion_option=QKVFusionOption.NONE, + sdpa_backend=SDPABackend.CUDNN, + ) + needs_hopper = any( + config is not None + and ( + config.sdpa_backend is not SDPABackend.CUDNN + or config.quantization.projection is not None + ) + for config in (self_config, cross_config) + ) + return AttentionBenchmarkCase( + implementation=( + f"self_{_implementation_id(self_config)}_" + f"cross_{_implementation_id(cross_config)}" + ), + self_attention_backend=( + AttentionBackend.OMNIDREAMS + if self_config is None + else AttentionBackend.OPTIMIZED + ), + cross_attention_backend=( + AttentionBackend.OMNIDREAMS + if cross_config is None + else AttentionBackend.OPTIMIZED + ), + self_attn_optimized_impl_config=self_config or reference_config, + cross_attn_optimized_impl_config=cross_config or reference_config, + minimum_compute_capability=(9, 0) if needs_hopper else None, + ) + + +_MODULE_CASE_MATRIX = tuple( + _block_case(self_config, cross_config) + for self_config in _MODULE_SELF_ATTENTION_CONFIGS + for cross_config in _MODULE_CROSS_ATTENTION_CONFIGS +) +"""Every valid self- and cross-attention implementation combination.""" + + +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, + self_attn_optimized_impl_config=case.self_attn_optimized_impl_config, + cross_attn_optimized_impl_config=case.cross_attn_optimized_impl_config, + ) + + +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, + self_attn_optimized_impl_config=config.self_attn_optimized_impl_config, + cross_attn_optimized_impl_config=config.cross_attn_optimized_impl_config, + ) + + 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_CASE_MATRIX, 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" + + 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( + "optimized_impl_config", + _MODULE_SELF_ATTENTION_CONFIGS, + ids=_implementation_id, +) +@torch.inference_mode() +def test_self_attention_benchmark( + benchmark: BenchmarkFixture, + optimized_impl_config: OptimizedImplConfig | None, +) -> 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") + block_case = _block_case(optimized_impl_config, None) + skip_unsupported_device(block_case, device) + dtype = torch.bfloat16 + config = _module_config(block_case) + attention = _make_block(config, block_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" + + # 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( + "optimized_impl_config", + _MODULE_CROSS_ATTENTION_CONFIGS, + ids=_implementation_id, +) +@torch.inference_mode() +def test_cross_attention_benchmark( + benchmark: BenchmarkFixture, + optimized_impl_config: OptimizedImplConfig | None, +) -> None: + """Benchmark cross-attention including production text KV projection.""" + if not torch.cuda.is_bf16_supported(): + pytest.skip("Omnidreams cross-attention benchmark requires bfloat16 support") + + device = torch.device("cuda") + block_case = _block_case(None, optimized_impl_config) + skip_unsupported_device(block_case, device) + dtype = torch.bfloat16 + config = _module_config(block_case) + attention = _make_block(config, block_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, + ) + torch.cuda.synchronize() + + benchmark.group = "omnidreams-dit-cross-attention" + + def synchronized_forward() -> torch.Tensor: + cache = attention.compute_kv(context) + 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..037d4e796 --- /dev/null +++ b/integrations/omnidreams/benchmarks/test_network.py @@ -0,0 +1,409 @@ +# 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, + self_attn_optimized_impl_config=case.self_attn_optimized_impl_config, + cross_attn_optimized_impl_config=case.cross_attn_optimized_impl_config, + ) + network = CosmosDiTNetwork(config).to(device=device, dtype=dtype) + network.eval() + network.update_parameters_after_loading_checkpoint() + 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.self_attn_optimized_impl_config is case.self_attn_optimized_impl_config + for block in network.blocks + ) + assert all( + block.cross_attn_optimized_impl_config is case.cross_attn_optimized_impl_config + 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" + + # Repeated scheduler evaluations overwrite one production steady-state slot. + cache.before_update(benchmark_chunk_idx) + + 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) + + 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() + + 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_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" + + # Repeated scheduler evaluations overwrite one production steady-state slot. + cache.start(benchmark_chunk_idx) + + 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) + + 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..6c0b5c1bd --- /dev/null +++ b/integrations/omnidreams/benchmarks/test_pipeline.py @@ -0,0 +1,310 @@ +# 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 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, + ) + + +@torch.inference_mode() +def _run_full_pipeline_benchmark( + benchmark: BenchmarkFixture, + *, + case: AttentionBenchmarkCase, +) -> None: + """Run one DiT backend full-pipeline 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, + "self_attn_optimized_impl_config": case.self_attn_optimized_impl_config, + "cross_attn_optimized_impl_config": case.cross_attn_optimized_impl_config, + }, + # 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 + + 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.self_attn_optimized_impl_config + is case.self_attn_optimized_impl_config + ) + assert ( + network_config.cross_attn_optimized_impl_config + is case.cross_attn_optimized_impl_config + ) + + 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 + assert native_executor._uses_fp8_dit is ( + case.native_dit_backend == "fp8_kvcache_cudnn" + ) + assert native_executor._attention_backend == case.native_attention_backend + else: + assert native_selection is None + assert native_executor is None + + benchmark.group = "omnidreams-full-pipeline-generate" + + next_chunk_index = cache_prefill_chunks + latest_output: torch.Tensor | None = None + + 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 + pipeline.finalize( + autoregressive_index=next_chunk_index, + cache=cache, + ) + torch.cuda.synchronize() + next_chunk_index += 1 + + output = benchmark.pedantic( + synchronized_generate, + teardown=teardown_generate, + iterations=1, + rounds=_BENCHMARK_ROUNDS, + warmup_rounds=_WARMUP_ROUNDS, + ) + + 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..0f97a828b 100644 --- a/integrations/omnidreams/omnidreams/config.py +++ b/integrations/omnidreams/omnidreams/config.py @@ -37,6 +37,7 @@ ) from omnidreams.runner import OmnidreamsRunnerConfig from omnidreams.transformer import CosmosTransformerConfig +from omnidreams.transformer.impl.modules import AttentionBackend from omnidreams.transformer.impl.network import ( CosmosDiTNetworkConfig, ) @@ -44,6 +45,11 @@ OmnidreamsWanVAEEncoderConfig as WanVAEEncoderConfig, ) +from flashdreams.accelerated.multi_head_attention.optimized import ( + QKVFusionOption, + SDPABackend, + OptimizedImplConfig, +) from flashdreams.infra.config import derive_config from flashdreams.infra.diffusion.model import DiffusionModelConfig from flashdreams.infra.diffusion.scheduler.fm import ( @@ -153,6 +159,68 @@ 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_TRITON_FA2 = cast( + OmnidreamsPipelineConfig, + derive_config( + SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE_PERF, + name=("omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae-triton-rtx-pro-6000"), + diffusion_model=dict( + transformer=dict( + network=dict( + self_attention_backend=AttentionBackend.OPTIMIZED, + cross_attention_backend=AttentionBackend.OPTIMIZED, + self_attn_optimized_impl_config=OptimizedImplConfig( + qkv_fusion_option=QKVFusionOption.FULL, + sdpa_backend=SDPABackend.FA2, + ), + cross_attn_optimized_impl_config=OptimizedImplConfig( + qkv_fusion_option=QKVFusionOption.NONE, + sdpa_backend=SDPABackend.FA2, + ), + ), + ), + ), + ), +) # ty:ignore[redundant-cast] +"""RTX Pro 6000 variant with Triton FA2 attention and FP8 projections.""" + +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 +480,10 @@ 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_TRITON_FA2, + 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, @@ -461,7 +533,50 @@ def _lightvae_fp8_state_path() -> str | None: description=( "Single-view chunk2 perf preset (compile + CUDA graphs across all stages)." ), - pipeline=SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE_PERF, + pipeline=derive_config( + SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE_PERF, + diffusion_model=dict(transformer=dict(skip_finalize_kv_cache=True)), + ), + prompt=_DEFAULT_PROMPT_1V, +) + +RUNNER_SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE_TRITON_FA2 = OmnidreamsRunnerConfig( + runner_name="omnidreams-triton-fa2", + description="Single-view chunk2 Triton FA2 preset tuned for RTX PRO 6000.", + pipeline=derive_config( + SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE_TRITON_FA2, + diffusion_model=dict(transformer=dict(skip_finalize_kv_cache=True)), + ), + 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=derive_config( + SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE_CUDA_CUDNN, + diffusion_model=dict(transformer=dict(skip_finalize_kv_cache=True)), + ), + 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=derive_config( + SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE_CUDA_SPARGE, + diffusion_model=dict(transformer=dict(skip_finalize_kv_cache=True)), + ), + 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=derive_config( + SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE_CUDA_SAGE3_FP8, + diffusion_model=dict(transformer=dict(skip_finalize_kv_cache=True)), + ), prompt=_DEFAULT_PROMPT_1V, ) @@ -569,6 +684,10 @@ 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_TRITON_FA2, + 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..81cf10510 100644 --- a/integrations/omnidreams/omnidreams/transformer/impl/modules.py +++ b/integrations/omnidreams/omnidreams/transformer/impl/modules.py @@ -17,6 +17,7 @@ import math from dataclasses import dataclass +from enum import Enum from typing import Literal import torch @@ -25,10 +26,34 @@ from torch import Tensor from torch.distributed import ProcessGroup +from flashdreams.accelerated.multi_head_attention import ( + AttentionConfig, + AttentionType, + QKNormScope, + RoPEConfig, + RoPEScope, + RoPEStyle, +) +from flashdreams.accelerated.multi_head_attention.optimized import ( + QKVFusionOption, + SDPABackend, + OptimizedImplConfig, + OptimizedHultiHeadAttention, +) 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.""" + + OPTIMIZED = "optimized" + """Use optimized attention for the selected branch.""" + + class GPT2FeedForward(nn.Module): """GPT-2 style feed-forward network with GELU activation.""" @@ -282,6 +307,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 +332,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 +363,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 +416,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 +431,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 +469,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 +478,263 @@ def forward( return super().forward(x, kv_cache, rope_freqs=None, update_kv_cache=False) +class OptimizedCrossAttention(OptimizedHultiHeadAttention): + """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", + optimized_impl_config: OptimizedImplConfig = OptimizedImplConfig( + qkv_fusion_option=QKVFusionOption.FUSE_KV, + sdpa_backend=SDPABackend.FA2, + ), + ) -> None: + """Initialize bias-free optimized 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. + optimized_impl_config: optimized backend and projection-fusion policies. + """ + del cp_method + super().__init__( + attention_type=AttentionType.CROSS_ATTENTION, + attention_config=AttentionConfig( + query_dim=query_dim, + context_dim=context_dim, + n_heads=n_heads, + head_dim=head_dim, + qk_norm_eps=1e-6, + qk_norm_scope=QKNormScope.HEAD, + ), + optimized_impl_config=optimized_impl_config, + ) + assert self.attention_config.context_dim is not None + self.q_proj = nn.Linear( + self.attention_config.query_dim, + self.attention_config.inner_dim, + bias=False, + ) + self.k_proj = nn.Linear( + self.attention_config.context_dim, + self.attention_config.inner_dim, + bias=False, + ) + self.v_proj = nn.Linear( + self.attention_config.context_dim, + self.attention_config.inner_dim, + bias=False, + ) + self.output_proj = nn.Linear( + self.attention_config.inner_dim, + self.attention_config.query_dim, + bias=False, + ) + self.q_norm = nn.RMSNorm( + self.attention_config.head_dim, eps=self.attention_config.qk_norm_eps + ) + self.k_norm = nn.RMSNorm( + self.attention_config.head_dim, eps=self.attention_config.qk_norm_eps + ) + self._initialize_derived_weights() + + def set_context_parallel_group(self, cp_group: ProcessGroup | None) -> None: + """Reject context parallelism unsupported by Optimized 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 Optimized 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 OptimizedSelfAttention(OptimizedHultiHeadAttention): + """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", + optimized_impl_config: OptimizedImplConfig = OptimizedImplConfig( + qkv_fusion_option=QKVFusionOption.FULL, + sdpa_backend=SDPABackend.FA2, + ), + ) -> None: + """Initialize bias-free optimized 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. + optimized_impl_config: optimized backend and projection-fusion policies. + + 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( + "Optimized self-attention requires context_dim to equal query_dim; " + f"got {context_dim} and {query_dim}" + ) + super().__init__( + attention_type=AttentionType.SELF_ATTENTION, + attention_config=AttentionConfig( + query_dim=query_dim, + context_dim=context_dim, + n_heads=n_heads, + head_dim=head_dim, + qk_norm_eps=1e-6, + qk_norm_scope=QKNormScope.HEAD, + rope_config=RoPEConfig( + style=RoPEStyle.SPLIT, + scope=RoPEScope.BEFORE_KV_CACHE, + ), + ), + optimized_impl_config=optimized_impl_config, + ) + assert self.attention_config.context_dim is not None + self.q_proj = nn.Linear( + self.attention_config.query_dim, + self.attention_config.inner_dim, + bias=False, + ) + self.k_proj = nn.Linear( + self.attention_config.context_dim, + self.attention_config.inner_dim, + bias=False, + ) + self.v_proj = nn.Linear( + self.attention_config.context_dim, + self.attention_config.inner_dim, + bias=False, + ) + self.output_proj = nn.Linear( + self.attention_config.inner_dim, + self.attention_config.query_dim, + bias=False, + ) + self.q_norm = nn.RMSNorm( + self.attention_config.head_dim, eps=self.attention_config.qk_norm_eps + ) + self.k_norm = nn.RMSNorm( + self.attention_config.head_dim, eps=self.attention_config.qk_norm_eps + ) + self._initialize_derived_weights() + + def set_context_parallel_group(self, cp_group: ProcessGroup | None) -> None: + """Reject context parallelism unsupported by Optimized 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 Optimized 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 +762,69 @@ 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, + self_attn_optimized_impl_config: OptimizedImplConfig = OptimizedImplConfig( + qkv_fusion_option=QKVFusionOption.FULL, + sdpa_backend=SDPABackend.FA2, + ), + cross_attn_optimized_impl_config: OptimizedImplConfig = OptimizedImplConfig( + qkv_fusion_option=QKVFusionOption.FUSE_KV, + sdpa_backend=SDPABackend.FA2, + ), ) -> 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.self_attn_optimized_impl_config = self_attn_optimized_impl_config + self.cross_attn_optimized_impl_config = cross_attn_optimized_impl_config # 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 = OptimizedSelfAttention( + query_dim=x_dim, + context_dim=None, + n_heads=num_heads, + head_dim=x_dim // num_heads, + cp_method=cp_method, + optimized_impl_config=self.self_attn_optimized_impl_config, + ) + + 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 = OptimizedCrossAttention( + query_dim=x_dim, + context_dim=context_dim, + n_heads=num_heads, + head_dim=x_dim // num_heads, + cp_method=cp_method, + optimized_impl_config=self.cross_attn_optimized_impl_config, + ) # MLP self.layer_norm_mlp = nn.LayerNorm(x_dim, elementwise_affine=False, eps=1e-6) @@ -555,13 +865,23 @@ 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 = OptimizedCrossAttention( + query_dim=x_dim, + context_dim=x_dim, + n_heads=num_heads, + head_dim=x_dim // num_heads, + cp_method=cp_method, + optimized_impl_config=self.cross_attn_optimized_impl_config, + ) def set_context_parallel_group( self, @@ -598,7 +918,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,7 +926,7 @@ def initialize_cache( device=device, dtype=dtype, ), - cross_attn=self.cross_attn.initialize_cache(context), + cross_attn=self.cross_attn.compute_kv(context), ) def forward( diff --git a/integrations/omnidreams/omnidreams/transformer/impl/network.py b/integrations/omnidreams/omnidreams/transformer/impl/network.py index 307e88220..cdef092ff 100644 --- a/integrations/omnidreams/omnidreams/transformer/impl/network.py +++ b/integrations/omnidreams/omnidreams/transformer/impl/network.py @@ -24,6 +24,11 @@ from torch import Tensor from torch.distributed import ProcessGroup +from flashdreams.accelerated.multi_head_attention.optimized import ( + QKVFusionOption, + SDPABackend, + OptimizedImplConfig, +) from flashdreams.core.distributed.context_parallel import ( cat_outputs_cp, split_inputs_cp, @@ -31,6 +36,7 @@ from flashdreams.infra.config import InstantiateConfig from .modules import ( + AttentionBackend, Block, BlockCache, FinalLayer, @@ -119,6 +125,28 @@ 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.""" + + self_attn_optimized_impl_config: OptimizedImplConfig = field( + default_factory=lambda: OptimizedImplConfig( + qkv_fusion_option=QKVFusionOption.FULL, + sdpa_backend=SDPABackend.FA2, + ) + ) + """Optimized implementation policy used by accelerated self-attention.""" + + cross_attn_optimized_impl_config: OptimizedImplConfig = field( + default_factory=lambda: OptimizedImplConfig( + qkv_fusion_option=QKVFusionOption.FUSE_KV, + sdpa_backend=SDPABackend.FA2, + ) + ) + """Optimized implementation policy used by accelerated cross-attention.""" + view_condition_dim: int = 16 """Embedding dim for the per-view conditioning vector.""" @@ -132,6 +160,8 @@ class CosmosDiTNetwork(nn.Module): def __init__(self, config: CosmosDiTNetworkConfig): super().__init__() self.config = config + self.self_attn_optimized_impl_config = config.self_attn_optimized_impl_config + self.cross_attn_optimized_impl_config = config.cross_attn_optimized_impl_config # add 1 for the condition mask in_channels = config.in_channels + 1 @@ -177,6 +207,10 @@ 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, + self_attn_optimized_impl_config=self.self_attn_optimized_impl_config, + cross_attn_optimized_impl_config=self.cross_attn_optimized_impl_config, ) for _ in range(self.config.num_blocks) ] diff --git a/integrations/omnidreams/pyproject.toml b/integrations/omnidreams/pyproject.toml index 2e38c80fb..a2498c61e 100644 --- a/integrations/omnidreams/pyproject.toml +++ b/integrations/omnidreams/pyproject.toml @@ -119,6 +119,10 @@ 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-triton-fa2" = "omnidreams.config:RUNNER_SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE_TRITON_FA2" +"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/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..034cc1239 100644 --- a/integrations/omnidreams/tests/test_recipe_configs.py +++ b/integrations/omnidreams/tests/test_recipe_configs.py @@ -33,6 +33,7 @@ import tomli as tomllib from omnidreams import config as config_mod from omnidreams.config import OMNIDREAMS_RUNNERS +from omnidreams.transformer import CosmosTransformerConfig from flashdreams.infra.runner import RunnerConfig @@ -51,12 +52,42 @@ 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-triton-fa2": ( + "omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae-triton-rtx-pro-6000" + ), + "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 assert all(slug == cfg.runner_name for slug, cfg in OMNIDREAMS_RUNNERS.items()) +def test_accelerated_runners_skip_finalize_kv_cache() -> None: + """Only accelerated public runners skip the finalize cache refresh.""" + skipped = set() + for slug, cfg in OMNIDREAMS_RUNNERS.items(): + transformer = cfg.pipeline.diffusion_model.transformer + assert isinstance(transformer, CosmosTransformerConfig) + if transformer.skip_finalize_kv_cache: + skipped.add(slug) + + assert skipped == { + "omnidreams-perf", + "omnidreams-triton-fa2", + "omnidreams-cuda-cudnn", + "omnidreams-cuda-sparge", + "omnidreams-cuda-sage3fp8", + } + + def test_runners_have_descriptions() -> None: """Every shipped runner needs a non-empty CLI description.""" empty = [ 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..626b972b3 --- /dev/null +++ b/integrations/omnidreams/tests/test_transformer_attention_backend.py @@ -0,0 +1,568 @@ +# 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.config import ( + OMNIDREAMS_CONFIGS, + SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE_TRITON_FA2, +) +from omnidreams.transformer import CosmosTransformerConfig +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.multi_head_attention import ( + AttentionType, +) +from flashdreams.accelerated.multi_head_attention import ( + optimized as optimized_attention, +) +from flashdreams.accelerated.multi_head_attention.optimized import ( + QKVFusionOption, + QuantizationOption, + SDPABackend, + OptimizedImplConfig, + OptimizedHultiHeadAttention, +) +from integrations.omnidreams.benchmarks.cases import BENCHMARK_CASES +from integrations.omnidreams.benchmarks.test_modules import ( + _MODULE_CASE_MATRIX, + _MODULE_CROSS_ATTENTION_CONFIGS, + _MODULE_SELF_ATTENTION_CONFIGS, + _OPTIMIZED_IMPL_CONFIGS, + _implementation_id, +) + +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) + + +def test_rtx_pro_6000_config_uses_triton_fa2_attention() -> None: + """Keep the RTX Pro 6000 preset on Triton FA2 attention.""" + config = SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE_TRITON_FA2 + assert config.name == ( + "omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae-triton-rtx-pro-6000" + ) + assert OMNIDREAMS_CONFIGS[config.name] is config + + transformer_config = config.diffusion_model.transformer + assert isinstance(transformer_config, CosmosTransformerConfig) + network_config = transformer_config.network + assert isinstance(network_config, CosmosDiTNetworkConfig) + assert network_config.self_attention_backend is AttentionBackend.OPTIMIZED + assert network_config.cross_attention_backend is AttentionBackend.OPTIMIZED + assert ( + network_config.self_attn_optimized_impl_config.sdpa_backend is SDPABackend.FA2 + ) + assert ( + network_config.cross_attn_optimized_impl_config.sdpa_backend is SDPABackend.FA2 + ) + assert ( + network_config.self_attn_optimized_impl_config.qkv_fusion_option + is QKVFusionOption.FULL + ) + assert ( + network_config.cross_attn_optimized_impl_config.qkv_fusion_option + is QKVFusionOption.NONE + ) + + +@pytest.mark.parametrize( + ("self_attention_backend", "cross_attention_backend"), + ( + (AttentionBackend.OPTIMIZED, AttentionBackend.OMNIDREAMS), + (AttentionBackend.OMNIDREAMS, AttentionBackend.OPTIMIZED), + ), + ids=("optimized-self", "optimized-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, + ) + + block = CosmosDiTNetwork(config).blocks[0] + + expected_self_type = ( + transformer_modules.OptimizedSelfAttention + if self_attention_backend is AttentionBackend.OPTIMIZED + else transformer_modules.SelfAttention + ) + expected_cross_type = ( + transformer_modules.OptimizedCrossAttention + if cross_attention_backend is AttentionBackend.OPTIMIZED + 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) + + +def test_optimized_attention_caches_cuda_capability( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Cache CUDA capability validation per device.""" + attention = transformer_modules.OptimizedSelfAttention( + query_dim=32, + n_heads=2, + head_dim=16, + ) + validated_devices: list[int] = [] + + monkeypatch.setattr(torch.cuda, "current_device", lambda: 3) + + def get_device_capability(device: int) -> tuple[int, int]: + validated_devices.append(device) + return (9, 0) + + monkeypatch.setattr(torch.cuda, "get_device_capability", get_device_capability) + + attention._validate_cuda_device("cuda") + attention._validate_cuda_device("cuda:3") + + assert validated_devices == [3] + + +@pytest.mark.parametrize( + "sdpa_backend", tuple(SDPABackend), ids=lambda backend: backend.value +) +def test_network_config_selects_optimized_attention( + sdpa_backend: SDPABackend, +) -> None: + """Propagate the configured self-attention SDPA implementation.""" + optimized_impl_config = OptimizedImplConfig( + qkv_fusion_option=QKVFusionOption.FULL, + sdpa_backend=sdpa_backend, + use_tma=False, + ) + 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.OPTIMIZED, + cross_attention_backend=AttentionBackend.OPTIMIZED, + self_attn_optimized_impl_config=optimized_impl_config, + ) + + network = CosmosDiTNetwork(config) + block = network.blocks[0] + + assert config.self_attn_optimized_impl_config is optimized_impl_config + assert network.self_attn_optimized_impl_config is optimized_impl_config + assert block.self_attn_optimized_impl_config is optimized_impl_config + self_attention = block.self_attn + assert isinstance(self_attention, OptimizedHultiHeadAttention) + assert self_attention.attention_type is AttentionType.SELF_ATTENTION + assert self_attention.optimized_impl_config is optimized_impl_config + assert self_attention.qkv_fusion_option is QKVFusionOption.FULL + assert self_attention.sdpa_backend is sdpa_backend + assert self_attention.use_tma is False + assert isinstance(self_attention.fused_qkv, torch.nn.Linear) + assert isinstance(self_attention.fused_kv, torch.nn.Linear) + 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.bfloat16 + assert cache._k.is_contiguous() + assert cache._v.is_contiguous() + assert isinstance(block.cross_attn, transformer_modules.OptimizedCrossAttention) + assert isinstance( + block.cross_view_attn, transformer_modules.OptimizedCrossAttention + ) + assert block.cross_attn.attention_config.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.sdpa_backend is SDPABackend.FA2 + assert block.cross_view_attn.attention_config.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.sdpa_backend is SDPABackend.FA2 + + +def test_network_config_selects_optimized_attention_policies() -> None: + """Propagate cross-attention SDPA and QKV fusion policies.""" + self_optimized_impl_config = OptimizedImplConfig( + qkv_fusion_option=QKVFusionOption.NONE, + sdpa_backend=SDPABackend.FA2, + use_tma=False, + ) + cross_optimized_impl_config = OptimizedImplConfig( + qkv_fusion_option=QKVFusionOption.NONE, + sdpa_backend=SDPABackend.CUDNN, + use_tma=False, + ) + + 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.OPTIMIZED, + cross_attention_backend=AttentionBackend.OPTIMIZED, + self_attn_optimized_impl_config=self_optimized_impl_config, + cross_attn_optimized_impl_config=cross_optimized_impl_config, + ) + + 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, OptimizedHultiHeadAttention) + assert isinstance(cross_attention, OptimizedHultiHeadAttention) + assert isinstance(cross_view_attention, OptimizedHultiHeadAttention) + + assert network.self_attn_optimized_impl_config is self_optimized_impl_config + assert network.cross_attn_optimized_impl_config is cross_optimized_impl_config + assert block.self_attn_optimized_impl_config is self_optimized_impl_config + assert block.cross_attn_optimized_impl_config is cross_optimized_impl_config + assert self_attention.optimized_impl_config is self_optimized_impl_config + assert cross_attention.optimized_impl_config is cross_optimized_impl_config + assert cross_view_attention.optimized_impl_config is cross_optimized_impl_config + assert not self_attention.use_tma and not cross_attention.use_tma + assert self_attention.sdpa_backend is SDPABackend.FA2 + 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 + + +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.self_attn_optimized_impl_config.sdpa_backend, + case.self_attn_optimized_impl_config.qkv_fusion_option, + case.cross_attn_optimized_impl_config.qkv_fusion_option, + ) + for case in pytorch_cases + ) == ( + ( + "omnidreams_torch", + AttentionBackend.OMNIDREAMS, + AttentionBackend.OMNIDREAMS, + SDPABackend.CUDNN, + QKVFusionOption.NONE, + QKVFusionOption.NONE, + ), + ( + "optimized_cudnn_fp8_self_full_no_tma_cross_none_tma", + AttentionBackend.OPTIMIZED, + AttentionBackend.OPTIMIZED, + SDPABackend.CUDNN, + QKVFusionOption.FULL, + QKVFusionOption.NONE, + ), + ( + "optimized_fa2_quantized_sdpa_self_full_tma_cross_none_tma", + AttentionBackend.OPTIMIZED, + AttentionBackend.OPTIMIZED, + SDPABackend.FA2, + QKVFusionOption.FULL, + QKVFusionOption.NONE, + ), + ) + gb300_best = pytorch_cases[1] + assert gb300_best.self_attn_optimized_impl_config == OptimizedImplConfig( + qkv_fusion_option=QKVFusionOption.FULL, + sdpa_backend=SDPABackend.CUDNN, + use_tma=False, + quantization=QuantizationOption(projection=torch.float8_e4m3fn), + ) + assert gb300_best.cross_attn_optimized_impl_config == OptimizedImplConfig( + qkv_fusion_option=QKVFusionOption.NONE, + sdpa_backend=SDPABackend.CUDNN, + use_tma=True, + quantization=QuantizationOption(projection=torch.float8_e4m3fn), + ) + + rtx_pro_6000_best = pytorch_cases[2] + assert rtx_pro_6000_best.self_attn_optimized_impl_config == OptimizedImplConfig( + qkv_fusion_option=QKVFusionOption.FULL, + sdpa_backend=SDPABackend.FA2, + use_tma=True, + quantization=QuantizationOption( + projection=torch.float8_e4m3fn, + quantized_sdpa=True, + ), + ) + assert rtx_pro_6000_best.cross_attn_optimized_impl_config == OptimizedImplConfig( + qkv_fusion_option=QKVFusionOption.NONE, + sdpa_backend=SDPABackend.FA2, + use_tma=True, + quantization=QuantizationOption( + projection=torch.float8_e4m3fn, + quantized_sdpa=True, + ), + ) + assert rtx_pro_6000_best.minimum_compute_capability == (9, 0) + + 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", (12, 0)), + ("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, fusion, TMA, and quantization policy.""" + expected_configs = { + OptimizedImplConfig( + sdpa_backend=sdpa_backend, + qkv_fusion_option=qkv_fusion_option, + quantization=QuantizationOption(projection=projection_dtype), + use_tma=use_tma, + ) + for sdpa_backend in SDPABackend + for qkv_fusion_option in QKVFusionOption + for use_tma in (False, True) + for projection_dtype in (None, torch.float8_e4m3fn) + } + expected_configs.update( + OptimizedImplConfig( + sdpa_backend=sdpa_backend, + qkv_fusion_option=qkv_fusion_option, + quantization=QuantizationOption( + projection=torch.float8_e4m3fn, + quantized_sdpa=True, + ), + use_tma=use_tma, + ) + for sdpa_backend in SDPABackend + for qkv_fusion_option in QKVFusionOption + for use_tma in (False, True) + ) + + assert set(_OPTIMIZED_IMPL_CONFIGS) == expected_configs + assert _MODULE_SELF_ATTENTION_CONFIGS[0] is None + assert set(_MODULE_SELF_ATTENTION_CONFIGS[1:]) == expected_configs + + expected_cross_configs = { + config + for config in expected_configs + if config.qkv_fusion_option is not QKVFusionOption.FULL + } + assert _MODULE_CROSS_ATTENTION_CONFIGS[0] is None + assert set(_MODULE_CROSS_ATTENTION_CONFIGS[1:]) == expected_cross_configs + + expected_block_ids = { + f"self_{_implementation_id(self_config)}_" + f"cross_{_implementation_id(cross_config)}" + for self_config in _MODULE_SELF_ATTENTION_CONFIGS + for cross_config in _MODULE_CROSS_ATTENTION_CONFIGS + } + assert {case.implementation for case in _MODULE_CASE_MATRIX} == expected_block_ids + assert len(_MODULE_SELF_ATTENTION_CONFIGS) == 37 + assert len(_MODULE_CROSS_ATTENTION_CONFIGS) == 25 + assert len(_MODULE_CASE_MATRIX) == 37 * 25 + assert len({case.pytest_id for case in _MODULE_CASE_MATRIX}) == len( + _MODULE_CASE_MATRIX + ) + + +def test_optimized_backend_preserves_checkpoint_keys() -> None: + """Load Omnidreams weights into the optimized block strictly.""" + omnidreams_block = Block(x_dim=32, context_dim=16, num_heads=2) + optimized_block = Block( + x_dim=32, + context_dim=16, + num_heads=2, + self_attention_backend=AttentionBackend.OPTIMIZED, + cross_attention_backend=AttentionBackend.OPTIMIZED, + ) + + optimized_block.load_state_dict(omnidreams_block.state_dict(), strict=True) + + +def test_optimized_cross_attention_dispatches_tma( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Keep the cross-attention adapter on the optimized backend-owned forward.""" + torch.manual_seed(0) + optimized_block = Block( + x_dim=32, + context_dim=16, + num_heads=2, + self_attention_backend=AttentionBackend.OPTIMIZED, + cross_attention_backend=AttentionBackend.OPTIMIZED, + ) + optimized_cross_attention = optimized_block.cross_attn + assert isinstance( + optimized_cross_attention, transformer_modules.OptimizedCrossAttention + ) + assert optimized_cross_attention.attention_type is AttentionType.CROSS_ATTENTION + assert ( + type(optimized_cross_attention).forward is OptimizedHultiHeadAttention.forward + ) + assert ( + type(optimized_cross_attention).compute_kv + is OptimizedHultiHeadAttention.compute_kv + ) + assert ( + type(optimized_cross_attention)._attention + is OptimizedHultiHeadAttention._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( + optimized_attention, + "is_tma_flash_attention_supported", + lambda *_: True, + ) + monkeypatch.setattr( + optimized_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) + optimized_output = optimized_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 optimized_output.shape == query.shape + assert torch.isfinite(optimized_output).all() diff --git a/pyproject.toml b/pyproject.toml index 02f62aa87..15ff95c45 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/multi_head_attention/triton/**", "flashdreams/flashdreams/accelerated/quantization/triton.py"] +ignore = ["flashdreams/flashdreams/accelerated/multi_head_attention/triton/**", "flashdreams/flashdreams/accelerated/quantization/triton.py"] 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/multi_head_attention/plot.py b/scripts/benchmark/flashdreams/accelerated/multi_head_attention/plot.py new file mode 100644 index 000000000..396256a29 --- /dev/null +++ b/scripts/benchmark/flashdreams/accelerated/multi_head_attention/plot.py @@ -0,0 +1,396 @@ +# 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 +import re +from pathlib import Path + +import matplotlib.pyplot as plt +from matplotlib.colors import CenteredNorm +from matplotlib.patches import Rectangle + +_DEFAULT_OUTPUT_DIR = Path( + "artifacts/benchmark/flashdreams/accelerated/multi_head_attention" +) +_DEFAULT_INPUT = _DEFAULT_OUTPUT_DIR / "benchmark.json" +_FASTEST_IMPLEMENTATION_COLUMN = "fastest-implementation-config" +_HIGHLIGHTED_ROWS = { + ("head", "before_kv_cache", False, False): "Cosmos", + ("inner", "before_kv_cache", True, True): "Wan", +} +_REFERENCE_IMPLEMENTATION = "reference-torch" + +RowKey = tuple[str, 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] = {} + + for record in payload["benchmarks"]: + if not isinstance(record, dict): + continue + group = record.get("group") + param = record.get("param") + if not isinstance(group, str) or not isinstance(param, str): + continue + match = re.fullmatch( + r"multi-head-attention-(self|cross)-norm-(.+)-rope-" + r"(interleaved|split)(?:-scope-(before-kv-cache|after-kv-cache))?" + r"-bias-(on|off)", + group, + ) + if match is None: + continue + attention, norm, rope, scope_state, bias_state = match.groups() + attention_type = f"{attention}_attention" + rope_scope = (scope_state or "before-kv-cache").replace("-", "_") + rope_interleaved = rope == "interleaved" + bias = bias_state == "on" + shared_id = group.removeprefix(f"multi-head-attention-{attention}-") + implementation_prefix = f"{shared_id}-" + if not param.startswith(implementation_prefix): + raise SystemExit(f"Unsupported attention parameter {param!r}") + implementation = param.removeprefix(implementation_prefix) + + 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_scope, 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 not values_ms: + raise SystemExit( + f"Benchmark JSON {input_path} contains no accelerated attention results" + ) + + subtitle_parts = ["Median latency in ms (lower is faster)"] + 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_scope, rope_interleaved, bias = row + rope = "interleaved" if rope_interleaved else "split" + scope = rope_scope.removesuffix("_kv_cache").replace("_", " ") + label = f"norm {norm} | rope {rope} {scope} cache | bias {'on' if bias else 'off'}" + highlight = _HIGHLIGHTED_ROWS.get((norm, rope_scope, rope_interleaved, bias)) + return f"{label} | {highlight}" if highlight is not None else label + + +def _column_label(column: str) -> str: + if column == _FASTEST_IMPLEMENTATION_COLUMN: + return "fastest implementation config" + return ( + " ".join( + "fuse qkv" if token == "full" else token for token in column.split("-") + ) + .replace(" projection ", "\nprojection ") + .replace(" quantized sdpa", "\nquantized sdpa") + ) + + +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: Torch-reference latency in milliseconds; ``None`` marks a + missing reference. + is_reference: Whether the cell is the torch reference 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{(1 - value / reference) * 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() + if _REFERENCE_IMPLEMENTATION not in columns: + raise SystemExit("Benchmark results have no reference-torch implementation") + display_columns = [*columns, _FASTEST_IMPLEMENTATION_COLUMN] + fastest_implementations: dict[RowKey, str | None] = {} + matrix = [] + for row in rows: + reference = values_ms.get((row, _REFERENCE_IMPLEMENTATION)) + fastest = min( + ( + column + for column in columns + if column != _REFERENCE_IMPLEMENTATION and (row, column) in values_ms + ), + key=lambda column: values_ms[(row, column)], + default=None, + ) + fastest_implementations[row] = fastest + relative_values = [ + values_ms.get((row, column), math.nan) / reference + if reference is not None + else math.nan + for column in columns + ] + relative_values.append( + values_ms[(row, fastest)] / reference + if reference is not None and fastest is not None + else math.nan + ) + matrix.append(relative_values) + 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(display_columns) * 2.0), + max(default_height, len(rows) * 0.9), + ) + image = axes.imshow( + matrix, + aspect="auto", + cmap="RdYlGn_r", + norm=CenteredNorm(vcenter=1.0), + ) + axes.set_xticks( + range(len(display_columns)), + labels=[_column_label(column) for column in display_columns], + rotation=45, + ha="right", + rotation_mode="anchor", + ) + axes.axvline(len(columns) - 0.5, color="black", linewidth=1.5) + axes.set_yticks(range(len(rows)), labels=[_row_label(row) for row in rows]) + for row_index, (_, norm, rope_scope, rope_interleaved, bias) in enumerate(rows): + if (norm, rope_scope, rope_interleaved, bias) in _HIGHLIGHTED_ROWS: + axes.add_patch( + Rectangle( + (-0.5, row_index - 0.5), + len(display_columns), + 1, + fill=False, + edgecolor="black", + linewidth=2.5, + clip_on=False, + zorder=3, + ) + ) + axes.set_xlabel("Implementation configuration") + axes.set_ylabel("Attention configuration") + axes.set_title(f"{attention_label} performance\n{subtitle}") + for row_index, row in enumerate(rows): + reference = values_ms.get((row, _REFERENCE_IMPLEMENTATION)) + for column_index, column in enumerate(display_columns): + if column == _FASTEST_IMPLEMENTATION_COLUMN: + fastest = fastest_implementations[row] + if fastest is None: + value = None + label = "N/A" + else: + value = values_ms[(row, fastest)] + relative_label = _cell_label( + value, reference, is_reference=False + ).splitlines()[1] + implementation, backend, config = _column_label(fastest).split( + maxsplit=2 + ) + label = f"{implementation} {backend}\n{config}\n{relative_label}" + else: + value = values_ms.get((row, column)) + label = _cell_label( + value, + reference, + is_reference=column == _REFERENCE_IMPLEMENTATION, + ) + text_color = "black" + relative_value = matrix[row_index][column_index] + if math.isfinite(relative_value): + red, green, blue, _ = image.cmap(image.norm(relative_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, + label, + ha="center", + va="center", + color=text_color, + ) + colorbar = figure.colorbar( + image, ax=axes, label="Runtime relative to torch reference (×)" + ) + 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/multi_head_attention/run.sh b/scripts/benchmark/flashdreams/accelerated/multi_head_attention/run.sh new file mode 100755 index 000000000..c50738fe8 --- /dev/null +++ b/scripts/benchmark/flashdreams/accelerated/multi_head_attention/run.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 + +cd "$(cd "$(dirname "${BASH_SOURCE[0]}")/../../../../.." && pwd)" +mkdir -p artifacts/benchmark/flashdreams/accelerated/multi_head_attention + +uv run --project flashdreams --group test pytest \ + flashdreams/benchmarks/accelerated/multi_head_attention \ + -p no:manual_marker -m manual --benchmark-only -v "$@" \ + --benchmark-json=artifacts/benchmark/flashdreams/accelerated/multi_head_attention/benchmark.json + +uv run python scripts/benchmark/flashdreams/accelerated/multi_head_attention/plot.py diff --git a/scripts/benchmark/flashdreams/accelerated/quantization/plot_gemm.py b/scripts/benchmark/flashdreams/accelerated/quantization/plot_gemm.py new file mode 100644 index 000000000..65296550e --- /dev/null +++ b/scripts/benchmark/flashdreams/accelerated/quantization/plot_gemm.py @@ -0,0 +1,373 @@ +# 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 quantized GEMM benchmark results.""" + +from __future__ import annotations + +import argparse +import json +import math +import re +from pathlib import Path + +import matplotlib.pyplot as plt +from matplotlib.colors import CenteredNorm + +_DEFAULT_OUTPUT_DIR = Path("artifacts/benchmark/flashdreams/accelerated/quantization") +_DEFAULT_INPUT = _DEFAULT_OUTPUT_DIR / "benchmark.json" +_FASTEST_CONFIG_COLUMN = "fastest-quantized-config" +_REFERENCE_CONFIG = "full-precision" +"""Benchmark parameter ID used as each row's performance reference.""" + +_CONFIG_LABELS = { + "full-precision": "Full precision", + "float8_e4m3fn-slice": "FP8 E4M3\nslice", + "float8_e4m3fn-tensor": "FP8 E4M3\ntensor", + "float8_e5m2-x-float8_e4m3fn-slice": "FP8 E5M2 × E4M3\nslice", + "float8_e5m2-x-float8_e4m3fn-tensor": "FP8 E5M2 × E4M3\ntensor", + "int8-slice": "INT8\nslice", + "int8-tensor": "INT8\ntensor", +} +_FORMAT_LABELS = {"fp16": "FP16", "bf16": "BF16", "fp32": "FP32"} +_SCOPE_LABELS = { + "end-to-end": "End-to-end", + "gemm-only": "GEMM only", +} + +RowKey = tuple[str, str] +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 quantized GEMM 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 quantized GEMM rows and median timings from benchmark JSON. + + Args: + input_path: Pytest-benchmark JSON file to parse. + + Returns: + Ordered row keys, configuration columns, median milliseconds by cell, + and plot subtitle. + + Raises: + SystemExit: The input cannot be read or does not contain compatible + quantized GEMM 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] = {} + + for record in payload["benchmarks"]: + if not isinstance(record, dict): + continue + group = record.get("group") + param = record.get("param") + if not isinstance(group, str) or not isinstance(param, str): + continue + match = re.fullmatch( + r"quantized-gemm-(fp16|bf16|fp32)-(end-to-end|gemm-only)", group + ) + if match is None: + continue + original_format, timing_scope = match.groups() + prefix = f"{original_format}-" + if not param.startswith(prefix): + raise SystemExit(f"Unsupported quantized GEMM parameter {param!r}") + configuration = param.removeprefix(prefix) + + 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 = (timing_scope, original_format) + cell = (row, configuration) + if cell in values_ms: + raise SystemExit(f"Duplicate benchmark cell for {row} and {configuration}") + if row not in rows: + rows.append(row) + if configuration not in columns: + columns.append(configuration) + values_ms[cell] = median * 1000.0 + + if not values_ms: + raise SystemExit( + f"Benchmark JSON {input_path} contains no quantized GEMM results" + ) + + subtitle_parts = ["Median latency in ms (lower is faster)"] + 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 _config_label(configuration: str) -> str: + """Return a compact display label for a benchmark configuration.""" + if configuration == _FASTEST_CONFIG_COLUMN: + return "Fastest quantized\nconfig" + return _CONFIG_LABELS.get(configuration, configuration.replace("-", " ")) + + +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: Full-precision latency in milliseconds; ``None`` marks a + missing reference. + is_reference: Whether the cell is the full-precision row reference. + + 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{(1 - value / reference) * 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, + timing_scope: str, + rows: list[RowKey], + columns: list[str], + values_ms: dict[CellKey, float], + subtitle: str, +) -> None: + """Write one timing scope's median-latency heatmap as a PNG. + + Args: + output_path: Destination PNG file. + timing_scope: Timed region represented by every row. + rows: Ordered source-format rows. + columns: Ordered GEMM configuration columns. + values_ms: Median milliseconds keyed by row and configuration. + subtitle: Benchmark environment summary. + + Raises: + SystemExit: Results have no full-precision reference configuration. + """ + if _REFERENCE_CONFIG not in columns: + raise SystemExit("Benchmark results have no full-precision configuration") + + display_columns = [*columns, _FASTEST_CONFIG_COLUMN] + fastest_configurations: dict[RowKey, str | None] = {} + matrix = [] + for row in rows: + reference = values_ms.get((row, _REFERENCE_CONFIG)) + fastest = min( + ( + column + for column in columns + if column != _REFERENCE_CONFIG and (row, column) in values_ms + ), + key=lambda column: values_ms[(row, column)], + default=None, + ) + fastest_configurations[row] = fastest + relative_values = [ + values_ms.get((row, column), math.nan) / reference + if reference is not None + else math.nan + for column in columns + ] + relative_values.append( + values_ms[(row, fastest)] / reference + if reference is not None and fastest is not None + else math.nan + ) + matrix.append(relative_values) + + figure, axes = plt.subplots() + default_width, default_height = figure.get_size_inches() + figure.set_size_inches( + max(default_width, len(display_columns) * 1.8), + max(default_height, len(rows) * 1.2), + ) + image = axes.imshow( + matrix, + aspect="auto", + cmap="RdYlGn_r", + norm=CenteredNorm(vcenter=1.0), + ) + axes.set_xticks( + range(len(display_columns)), + labels=[_config_label(column) for column in display_columns], + rotation=45, + ha="right", + rotation_mode="anchor", + ) + axes.axvline(len(columns) - 0.5, color="black", linewidth=1.5) + axes.set_yticks( + range(len(rows)), + labels=[_FORMAT_LABELS.get(row[1], row[1]) for row in rows], + ) + axes.set_xlabel("GEMM configuration") + axes.set_ylabel("Source dtype") + axes.set_title(f"Quantized GEMM {_SCOPE_LABELS[timing_scope]}\n{subtitle}") + + for row_index, row in enumerate(rows): + reference = values_ms.get((row, _REFERENCE_CONFIG)) + for column_index, column in enumerate(display_columns): + if column == _FASTEST_CONFIG_COLUMN: + fastest = fastest_configurations[row] + if fastest is None: + value = None + label = "N/A" + else: + value = values_ms[(row, fastest)] + label = ( + f"{_config_label(fastest)}\n" + f"{_cell_label(value, reference, is_reference=False)}" + ) + else: + value = values_ms.get((row, column)) + label = _cell_label( + value, + reference, + is_reference=column == _REFERENCE_CONFIG, + ) + + text_color = "black" + relative_value = matrix[row_index][column_index] + if math.isfinite(relative_value): + red, green, blue, _ = image.cmap(image.norm(relative_value)) + linear = tuple( + channel / 12.92 + if channel <= 0.04045 + else ((channel + 0.055) / 1.055) ** 2.4 + for channel in (red, green, blue) + ) + 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, + label, + ha="center", + va="center", + color=text_color, + ) + + colorbar = figure.colorbar( + image, ax=axes, label="Runtime relative to full precision (×)" + ) + 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 end-to-end and GEMM-only performance matrices.""" + args = _parse_args(argv) + rows, columns, values_ms, subtitle = _load_matrix(args.input) + timing_scopes = list(dict.fromkeys(row[0] for row in rows)) + for timing_scope in timing_scopes: + panel_rows = [row for row in rows if row[0] == timing_scope] + panel_values = { + cell: value + for cell, value in values_ms.items() + if cell[0][0] == timing_scope + } + output = args.output_dir / f"{timing_scope.replace('-', '_')}.png" + _write_png( + output, + timing_scope, + panel_rows, + columns, + panel_values, + subtitle, + ) + print(f"Wrote {output}") + + +if __name__ == "__main__": + main() diff --git a/scripts/benchmark/flashdreams/accelerated/quantization/plot_quantized_linear.py b/scripts/benchmark/flashdreams/accelerated/quantization/plot_quantized_linear.py new file mode 100644 index 000000000..fc2370b0b --- /dev/null +++ b/scripts/benchmark/flashdreams/accelerated/quantization/plot_quantized_linear.py @@ -0,0 +1,415 @@ +# 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 quantized linear benchmark results.""" + +from __future__ import annotations + +import argparse +import json +import math +import re +from pathlib import Path + +import matplotlib.pyplot as plt +from matplotlib.colors import CenteredNorm + +_DEFAULT_OUTPUT_DIR = Path("artifacts/benchmark/flashdreams/accelerated/quantization") +_DEFAULT_INPUT = _DEFAULT_OUTPUT_DIR / "benchmark.json" +_FASTEST_CONFIG_COLUMN = "fastest-quantized-config" +_REFERENCE_CONFIG = "nn-linear" +"""Benchmark parameter ID used as each row's performance reference.""" + +_CONFIG_PATTERN = re.compile( + r"(?Pfloat8_e4m3fn|float8_e5m2-x-float8_e4m3fn|int8)-" + r"weight-(?Pper_out_channel|tensor)-" + r"input-(?Pslice|tensor)-" + r"(?Pfull-precision-x|prequantized-x)" +) +_QUANTIZED_FORMAT_LABELS = { + "float8_e4m3fn": "FP8 E4M3", + "float8_e5m2-x-float8_e4m3fn": "FP8 E5M2 × E4M3", + "int8": "INT8", +} +_INPUT_STATE_LABELS = { + "full-precision-x": "Quantize input", + "prequantized-x": "Prequantized input", +} +_INPUT_STATE_OUTPUT_NAMES = { + "full-precision-x": "quantize_input", + "prequantized-x": "prequantized_input", +} +_SOURCE_FORMAT_LABELS = {"fp16": "FP16", "bf16": "BF16", "fp32": "FP32"} + +CellKey = 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 quantized linear 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 _parse_configuration(configuration: str) -> re.Match[str]: + """Parse a quantized linear configuration ID. + + Args: + configuration: Pytest parameter suffix excluding the source dtype. + + Returns: + Match containing quantized format, weight and input granularities, and + input state. + + Raises: + SystemExit: The configuration ID is unsupported. + """ + match = _CONFIG_PATTERN.fullmatch(configuration) + if match is None: + raise SystemExit( + f"Unsupported quantized linear configuration {configuration!r}" + ) + return match + + +def _load_matrix( + input_path: Path, +) -> tuple[list[str], list[str], dict[CellKey, float], str]: + """Load quantized linear rows and median timings from benchmark JSON. + + Args: + input_path: Pytest-benchmark JSON file to parse. + + Returns: + Ordered source-format rows, configuration columns, median milliseconds + by cell, and plot subtitle. + + Raises: + SystemExit: The input cannot be read or does not contain compatible + quantized linear 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[str] = [] + columns: list[str] = [] + values_ms: dict[CellKey, float] = {} + + for record in payload["benchmarks"]: + if not isinstance(record, dict): + continue + group = record.get("group") + param = record.get("param") + if not isinstance(group, str) or not isinstance(param, str): + continue + match = re.fullmatch(r"quantized-linear-(fp16|bf16|fp32)", group) + if match is None: + continue + source_format = match.group(1) + prefix = f"{source_format}-" + if not param.startswith(prefix): + raise SystemExit(f"Unsupported quantized linear parameter {param!r}") + configuration = param.removeprefix(prefix) + if configuration != _REFERENCE_CONFIG: + _parse_configuration(configuration) + + 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" + ) + + cell = (source_format, configuration) + if cell in values_ms: + raise SystemExit( + f"Duplicate benchmark cell for {source_format} and {configuration}" + ) + if source_format not in rows: + rows.append(source_format) + if configuration not in columns: + columns.append(configuration) + values_ms[cell] = median * 1000.0 + + if not values_ms: + raise SystemExit( + f"Benchmark JSON {input_path} contains no quantized linear results" + ) + + subtitle_parts = ["Median latency in ms (lower is faster)"] + 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 _config_label(configuration: str) -> str: + """Return a compact display label for a benchmark configuration.""" + if configuration == _REFERENCE_CONFIG: + return "nn.Linear" + if configuration == _FASTEST_CONFIG_COLUMN: + return "Fastest quantized\nconfig" + match = _parse_configuration(configuration) + quantized_format = _QUANTIZED_FORMAT_LABELS[match.group("format")] + weight = ( + "weight per output channel" + if match.group("weight") == "per_out_channel" + else "weight tensor" + ) + return f"{quantized_format}\n{weight}\ninput {match.group('input')}" + + +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: ``nn.Linear`` latency in milliseconds; ``None`` marks a + missing reference. + is_reference: Whether the cell is the ``nn.Linear`` row reference. + + 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{(1 - value / reference) * 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, + input_state: str, + rows: list[str], + columns: list[str], + values_ms: dict[CellKey, float], + subtitle: str, +) -> None: + """Write one input state's median-latency heatmap as a PNG. + + Args: + output_path: Destination PNG file. + input_state: Whether input quantization occurs inside the timed region. + rows: Ordered source-format rows. + columns: Ordered linear configuration columns. + values_ms: Median milliseconds keyed by row and configuration. + subtitle: Benchmark environment summary. + + Raises: + SystemExit: Results have no complete ``nn.Linear`` reference column. + """ + if _REFERENCE_CONFIG not in columns or any( + (row, _REFERENCE_CONFIG) not in values_ms for row in rows + ): + raise SystemExit("Benchmark results have no complete nn.Linear reference") + + display_columns = [*columns, _FASTEST_CONFIG_COLUMN] + fastest_configurations: dict[str, str | None] = {} + matrix = [] + for row in rows: + reference = values_ms[(row, _REFERENCE_CONFIG)] + fastest = min( + ( + column + for column in columns + if column != _REFERENCE_CONFIG and (row, column) in values_ms + ), + key=lambda column: values_ms[(row, column)], + default=None, + ) + fastest_configurations[row] = fastest + relative_values = [ + values_ms.get((row, column), math.nan) / reference for column in columns + ] + relative_values.append( + values_ms[(row, fastest)] / reference if fastest is not None else math.nan + ) + matrix.append(relative_values) + + 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(display_columns) * 1.9), + max(default_height, len(rows) * 2.0), + ) + image = axes.imshow( + matrix, + aspect="auto", + cmap="RdYlGn_r", + norm=CenteredNorm(vcenter=1.0), + ) + axes.set_xticks( + range(len(display_columns)), + labels=[_config_label(column) for column in display_columns], + rotation=45, + ha="right", + rotation_mode="anchor", + ) + axes.axvline(len(columns) - 0.5, color="black", linewidth=1.5) + axes.set_yticks( + range(len(rows)), + labels=[_SOURCE_FORMAT_LABELS.get(row, row) for row in rows], + ) + axes.set_xlabel("Linear configuration") + axes.set_ylabel("Source dtype") + axes.set_title(f"Quantized linear · {_INPUT_STATE_LABELS[input_state]}\n{subtitle}") + + for row_index, row in enumerate(rows): + reference = values_ms[(row, _REFERENCE_CONFIG)] + for column_index, column in enumerate(display_columns): + if column == _FASTEST_CONFIG_COLUMN: + fastest = fastest_configurations[row] + if fastest is None: + value = None + label = "N/A" + else: + value = values_ms[(row, fastest)] + label = ( + f"{_config_label(fastest)}\n" + f"{_cell_label(value, reference, is_reference=False)}" + ) + else: + value = values_ms.get((row, column)) + label = _cell_label( + value, + reference, + is_reference=column == _REFERENCE_CONFIG, + ) + + text_color = "black" + relative_value = matrix[row_index][column_index] + if math.isfinite(relative_value): + red, green, blue, _ = image.cmap(image.norm(relative_value)) + linear = tuple( + channel / 12.92 + if channel <= 0.04045 + else ((channel + 0.055) / 1.055) ** 2.4 + for channel in (red, green, blue) + ) + 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, + label, + ha="center", + va="center", + color=text_color, + ) + + colorbar = figure.colorbar( + image, ax=axes, label="Runtime relative to nn.Linear (×)" + ) + 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 quantize-input and prequantized-input matrices.""" + args = _parse_args(argv) + rows, columns, values_ms, subtitle = _load_matrix(args.input) + input_states = list( + dict.fromkeys( + _parse_configuration(column).group("input_state") + for column in columns + if column != _REFERENCE_CONFIG + ) + ) + for input_state in input_states: + panel_columns = [ + column + for column in columns + if column == _REFERENCE_CONFIG + or _parse_configuration(column).group("input_state") == input_state + ] + output_name = _INPUT_STATE_OUTPUT_NAMES[input_state] + output = args.output_dir / f"quantized_linear_{output_name}.png" + _write_png( + output, + input_state, + rows, + panel_columns, + values_ms, + subtitle, + ) + print(f"Wrote {output}") + + +if __name__ == "__main__": + main() diff --git a/scripts/benchmark/flashdreams/accelerated/quantization/plot_quantizer.py b/scripts/benchmark/flashdreams/accelerated/quantization/plot_quantizer.py new file mode 100644 index 000000000..237dbe1be --- /dev/null +++ b/scripts/benchmark/flashdreams/accelerated/quantization/plot_quantizer.py @@ -0,0 +1,331 @@ +# 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 Torch and Triton quantizer benchmarks.""" + +from __future__ import annotations + +import argparse +import json +import math +import re +from pathlib import Path + +import matplotlib.pyplot as plt +import numpy as np +from matplotlib.colors import CenteredNorm + +_DEFAULT_OUTPUT_DIR = Path("artifacts/benchmark/flashdreams/accelerated/quantization") +_DEFAULT_INPUT = _DEFAULT_OUTPUT_DIR / "benchmark.json" + +_OPERATIONS = ("quantize", "dequantize") +_FORMATS = ("float8_e4m3fn", "float8_e5m2", "int8") +_GRANULARITIES = ("slice", "tensor") +_IMPLEMENTATIONS = ("torch", "triton") +_COLUMNS = tuple( + (granularity, implementation) + for granularity in _GRANULARITIES + for implementation in _IMPLEMENTATIONS +) +_GROUP_PATTERN = re.compile( + r"(?Pquantize|dequantize)-" + r"(?Pfloat8_e4m3fn|float8_e5m2|int8)-" + r"(?Pslice|tensor)" +) +_FORMAT_LABELS = { + "float8_e4m3fn": "FP8 E4M3", + "float8_e5m2": "FP8 E5M2", + "int8": "INT8", +} + +CellKey = tuple[str, str, 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 Torch and Triton quantizer 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_results(input_path: Path) -> tuple[dict[CellKey, float], str]: + """Load quantizer median timings from benchmark JSON. + + Args: + input_path: Pytest-benchmark JSON file to parse. + + Returns: + Median milliseconds by operation, format, granularity, and + implementation, plus the plot subtitle. + + Raises: + SystemExit: The input cannot be read or lacks complete quantizer 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_ms: dict[CellKey, float] = {} + for record in payload["benchmarks"]: + if not isinstance(record, dict): + continue + group = record.get("group") + match = _GROUP_PATTERN.fullmatch(group) if isinstance(group, str) else None + if match is None: + continue + + extra_info = record.get("extra_info") + implementation = ( + extra_info.get("implementation") if isinstance(extra_info, dict) else None + ) + if implementation not in _IMPLEMENTATIONS: + raise SystemExit( + f"Benchmark {record.get('name', '')} has no supported " + "implementation metadata" + ) + + 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" + ) + + cell = ( + match.group("operation"), + match.group("format"), + match.group("granularity"), + implementation, + ) + if cell in values_ms: + raise SystemExit(f"Duplicate quantizer benchmark cell for {cell}") + values_ms[cell] = median * 1000.0 + + missing_operations = [ + operation + for operation in _OPERATIONS + if not any(cell[0] == operation for cell in values_ms) + ] + if missing_operations: + raise SystemExit( + f"Benchmark JSON {input_path} lacks quantizer operations: " + f"{', '.join(missing_operations)}" + ) + return values_ms, _subtitle(payload) + + +def _subtitle(payload: dict[str, object]) -> str: + """Build a compact benchmark-environment subtitle.""" + parts = ["Median latency in ms (lower is faster)"] + 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 _cell_label( + value: float | None, reference: float | None, *, is_reference: bool +) -> str: + """Format a latency and its relationship to the Torch reference. + + Args: + value: Cell latency in milliseconds; ``None`` marks a missing result. + reference: Torch latency for the same format and granularity. + is_reference: Whether the cell is the Torch reference. + + Returns: + Two-line latency and relative-performance annotation. + """ + if value is None: + return "N/A" + if is_reference: + return f"{value:.3f} ms\n1.00× reference" + if reference is None: + return f"{value:.3f} ms\nreference unavailable" + if value < reference: + return f"{value:.3f} ms\n{reference / value:.2f}× faster" + if value > reference: + return f"{value:.3f} ms\n{value / reference:.2f}× slower" + return f"{value:.3f} ms\nsame as reference" + + +def _write_png( + output_path: Path, + operation: str, + values_ms: dict[CellKey, float], + subtitle: str, +) -> None: + """Write one operation's Torch-versus-Triton latency heatmap. + + Args: + output_path: Destination PNG file. + operation: Quantizer operation represented by the matrix. + values_ms: Median milliseconds keyed by benchmark configuration. + subtitle: Benchmark environment summary. + + Raises: + SystemExit: Any format and granularity pair lacks its Torch reference. + """ + missing_references = [ + (format, granularity) + for format in _FORMATS + for granularity in _GRANULARITIES + if (operation, format, granularity, "torch") not in values_ms + ] + if missing_references: + raise SystemExit( + f"{operation.capitalize()} results lack Torch references: " + + ", ".join( + f"{format}/{granularity}" for format, granularity in missing_references + ) + ) + + matrix: list[list[float]] = [] + for format in _FORMATS: + row = [] + for granularity, implementation in _COLUMNS: + reference = values_ms[(operation, format, granularity, "torch")] + value = values_ms.get((operation, format, granularity, implementation)) + row.append(value / reference if value is not None else math.nan) + matrix.append(row) + + figure, axes = plt.subplots() + default_width, default_height = figure.get_size_inches() + figure.set_size_inches( + max(default_width, len(_COLUMNS) * 1.8), + max(default_height, len(_FORMATS) * 1.2), + ) + image = axes.imshow( + np.asarray(matrix), + aspect="auto", + cmap="RdYlGn_r", + norm=CenteredNorm(vcenter=1.0), + ) + axes.set_xticks( + range(len(_COLUMNS)), + labels=[ + f"{granularity.capitalize()}\n{implementation.capitalize()}" + for granularity, implementation in _COLUMNS + ], + ) + axes.set_yticks( + range(len(_FORMATS)), + labels=[_FORMAT_LABELS[format] for format in _FORMATS], + ) + axes.set_xlabel("Scale granularity and implementation") + axes.set_ylabel("Quantized dtype") + axes.set_title(f"{operation.capitalize()} · Torch vs Triton\n{subtitle}") + + for row_index, format in enumerate(_FORMATS): + for column_index, (granularity, implementation) in enumerate(_COLUMNS): + reference = values_ms[(operation, format, granularity, "torch")] + value = values_ms.get((operation, format, granularity, implementation)) + label = _cell_label( + value, + reference, + is_reference=implementation == "torch", + ) + + text_color = "black" + relative_value = matrix[row_index][column_index] + if math.isfinite(relative_value): + colors = image.cmap(image.norm(np.asarray([relative_value]))) + red, green, blue, _ = np.asarray(colors)[0] + linear = tuple( + channel / 12.92 + if channel <= 0.04045 + else ((channel + 0.055) / 1.055) ** 2.4 + for channel in (red, green, blue) + ) + 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, + label, + ha="center", + va="center", + color=text_color, + ) + + colorbar = figure.colorbar(image, ax=axes, label="Runtime relative to Torch (×)") + 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 quantize and dequantize performance matrices.""" + args = _parse_args(argv) + values_ms, subtitle = _load_results(args.input) + for operation in _OPERATIONS: + output = args.output_dir / f"{operation}.png" + _write_png(output, operation, values_ms, subtitle) + print(f"Wrote {output}") + + +if __name__ == "__main__": + main() diff --git a/scripts/benchmark/flashdreams/accelerated/quantization/run.sh b/scripts/benchmark/flashdreams/accelerated/quantization/run.sh new file mode 100755 index 000000000..a9432ab5e --- /dev/null +++ b/scripts/benchmark/flashdreams/accelerated/quantization/run.sh @@ -0,0 +1,19 @@ +#!/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/quantization + +uv run --project flashdreams --group test pytest \ + flashdreams/benchmarks/accelerated/quantization/test_quantizer_benchmark.py \ + flashdreams/benchmarks/accelerated/quantization/test_quantized_gemm_benchmark.py \ + flashdreams/benchmarks/accelerated/quantization/test_quantized_linear_benchmark.py \ + -p no:manual_marker -m manual --benchmark-only -v "$@" \ + --benchmark-json=artifacts/benchmark/flashdreams/accelerated/quantization/benchmark.json + +uv run python scripts/benchmark/flashdreams/accelerated/quantization/plot_quantizer.py +uv run python scripts/benchmark/flashdreams/accelerated/quantization/plot_gemm.py +uv run python scripts/benchmark/flashdreams/accelerated/quantization/plot_quantized_linear.py diff --git a/scripts/benchmark/omnidreams/plot.py b/scripts/benchmark/omnidreams/plot.py new file mode 100644 index 000000000..8f713656b --- /dev/null +++ b/scripts/benchmark/omnidreams/plot.py @@ -0,0 +1,630 @@ +# 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 plots for Omnidreams benchmark results.""" + +from __future__ import annotations + +import argparse +import json +import math +import re +from pathlib import Path + +import matplotlib.pyplot as plt +import numpy as np +from matplotlib.colors import CenteredNorm + +_DEFAULT_INPUT = Path("artifacts/benchmark/omnidreams/benchmark.json") +_DEFAULT_OUTPUT_DIR = Path("artifacts/benchmark/omnidreams") + +_ATTENTION_PANELS = ( + ("omnidreams-dit-self-attention", "Self-attention"), + ("omnidreams-dit-cross-attention", "Cross-attention"), +) +_DIT_BLOCK_PANEL = ("omnidreams-dit-block", "DiT block") +_PIPELINE_GENERATE_PANEL = ( + "omnidreams-full-pipeline-generate", + "Pipeline generate", +) +_END_TO_END_PANELS = ( + ("omnidreams-dit-network", "Network eval"), + _PIPELINE_GENERATE_PANEL, +) +_PANELS = (*_ATTENTION_PANELS, _DIT_BLOCK_PANEL, *_END_TO_END_PANELS) + +_NATIVE_LABELS = { + "omnidreams-torch": "PyTorch Omnidreams BF16", + "cuda": "CUDA cuDNN FP8", + "cuda-sparge": "CUDA Sparge FP8", + "cuda-sage3": "CUDA Sage3 BF16", + "cuda-sage3-fp8": "CUDA Sage3 FP8", +} +"""Implementation labels for configurations without parseable optimized IDs.""" + +_FUSION_LABELS = { + "none": "None", + "fuse-kv": "Fuse KV", + "full": "Fuse QKV", +} +"""Display labels by stable QKV fusion ID.""" + +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 three figures." + ) + 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})", + ) + parser.add_argument( + "--pipeline-generate-only", + action="store_true", + help="write only the pipeline generate benchmark figure", + ) + 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] = [] + + 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 + param = record.get("param") + if not isinstance(param, str) or not param: + raise SystemExit( + f"Benchmark {record.get('name', '')} has no parameter ID" + ) + implementation = param + 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) + + missing = [group for group, results in values.items() if not results] + if missing: + raise SystemExit( + f"Benchmark JSON {input_path} lacks groups: {', '.join(missing)}" + ) + return values, configurations, _subtitle(payload) + + +def _subtitle(payload: dict[str, object]) -> str: + """Build a compact benchmark-environment subtitle.""" + parts = ["Median latency in ms (lower is faster)"] + 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 _backend_label(backend: str) -> str: + return "cuDNN" if backend == "cudnn" else backend.upper() + + +def _attention_configuration_label(implementation: str) -> str: + """Format one attention implementation as a compact axis label. + + Args: + implementation: Stable benchmark implementation identifier. + + Returns: + Backend, fusion, TMA, and projection-quantization label. + """ + if implementation == "omnidreams": + return "Omnidreams" + match = re.fullmatch( + r"(?:optimized|triton)-(cudnn|fa2)-(none|full|fuse-kv)-(no-tma|tma)" + r"(?:-projection-(float8-e4m3fn))?", + implementation, + ) + if match is None: + return implementation.replace("-", " ") + backend, fusion, tma, projection = match.groups() + tma_label = "no TMA" if tma == "no-tma" else "TMA" + projection_label = "" if projection is None else "\nProjection: FP8 E4M3" + return ( + f"Optimized {_backend_label(backend)}\n" + f"{_FUSION_LABELS[fusion]} · {tma_label}{projection_label}" + ) + + +def _end_to_end_configuration_label(implementation: str) -> str: + """Format one network or pipeline implementation as an axis label. + + Args: + implementation: Stable benchmark implementation identifier. + + Returns: + Implementation, self-attention, and cross-attention label. + """ + native_label = _NATIVE_LABELS.get(implementation) + if native_label is not None: + return native_label + + selected = re.fullmatch( + r"(?:optimized|triton)-(cudnn|fa2)-(bf16)-self-(none|full|fuse-kv)-" + r"(no-tma|tma)-cross-(none|fuse-kv)-(no-tma|tma)", + implementation, + ) + if selected is not None: + backend, precision, self_fusion, self_tma, cross_fusion, cross_tma = ( + selected.groups() + ) + self_tma_label = "no TMA" if self_tma == "no-tma" else "TMA" + cross_tma_label = "no TMA" if cross_tma == "no-tma" else "TMA" + return ( + f"Optimized {_backend_label(backend)} {precision.upper()}\n" + f"Self: {_FUSION_LABELS[self_fusion]} · {self_tma_label}\n" + f"Cross: {_FUSION_LABELS[cross_fusion]} · {cross_tma_label}" + ) + + legacy = re.fullmatch( + r"(?:optimized|triton)-(cudnn|fa2)-(bf16)-(none|full|fuse-kv)" + r"(-omnidreams-cross)?", + implementation, + ) + if legacy is None: + return implementation.replace("-", " ") + backend, precision, self_fusion, omnidreams_cross = legacy.groups() + cross_fusion = "fuse-kv" if self_fusion == "full" else self_fusion + cross_label = ( + "Omnidreams · None" + if omnidreams_cross is not None + else f"{_FUSION_LABELS[cross_fusion]} · TMA" + ) + return ( + f"Optimized {_backend_label(backend)} {precision.upper()}\n" + f"Self: {_FUSION_LABELS[self_fusion]} · TMA\n" + f"Cross: {cross_label}" + ) + + +def _bar_label( + latency_ms: float, + reference_ms: float, + *, + is_reference: bool, +) -> str: + """Format latency and relative performance against a reference. + + Args: + latency_ms: Bar latency in milliseconds. + reference_ms: Reference latency in milliseconds. + is_reference: Whether the bar or cell is the 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{(1 - latency_ms / reference_ms) * 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 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 = all(panel in _END_TO_END_PANELS for panel in panels) + reference_configuration = "omnidreams-torch" if end_to_end else "omnidreams" + configuration_label = ( + _end_to_end_configuration_label + if end_to_end + else _attention_configuration_label + ) + panel_configurations = _panel_configurations(panels, values, configurations) + figure, axes = plt.subplots( + len(panels), + 1, + figsize=( + max(16.0, len(panel_configurations) * 1.3), + len(panels) * 5.0, + ), + layout="constrained", + ) + axes_list = [axes] if len(panels) == 1 else list(axes) + for axes_item, (group, panel_title) in zip(axes_list, panels, strict=True): + reference = values[group].get(reference_configuration) + if reference is None: + raise SystemExit( + f"Benchmark results for {group} have no {reference_configuration} reference" + ) + 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 == reference_configuration + else "C1" + if configuration.startswith("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 == reference_configuration, + ), + 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) + for configuration in ordered_configurations + ], + rotation=45, + ha="right", + rotation_mode="anchor", + ) + + figure.supxlabel( + "Configuration (self-attention + cross-attention)" + if end_to_end + else "Configuration (SDPA backend / QKV fusion / TMA / projection dtype)" + ) + figure.suptitle(f"{title}\n{subtitle}") + output_path.parent.mkdir(parents=True, exist_ok=True) + figure.savefig(output_path) + plt.close(figure) + + +def _split_block_implementation( + implementation: str, +) -> tuple[str, str]: + """Split one DiT block ID into self- and cross-attention implementations. + + Args: + implementation: Stable DiT block benchmark identifier. + + Returns: + Self- and cross-attention implementation identifiers. + + Raises: + SystemExit: The identifier does not encode both implementations. + """ + prefix = "self-" + if not implementation.startswith(prefix): + raise SystemExit(f"Unsupported DiT block parameter {implementation!r}") + self_implementation, separator, cross_implementation = implementation[ + len(prefix) : + ].partition("-cross-") + if not separator or not self_implementation or not cross_implementation: + raise SystemExit(f"Unsupported DiT block parameter {implementation!r}") + return self_implementation, cross_implementation + + +def _write_block_figure( + output_path: Path, + values: BenchmarkValues, + subtitle: str, +) -> None: + """Write the DiT block self-by-cross implementation heatmap. + + Args: + output_path: Destination PNG file. + values: Median milliseconds by group and implementation. + subtitle: Benchmark environment summary. + """ + block_values = values[_DIT_BLOCK_PANEL[0]] + self_configurations: list[str] = [] + cross_configurations: list[str] = [] + latencies: dict[tuple[str, str], float] = {} + for implementation, latency in block_values.items(): + self_implementation, cross_implementation = _split_block_implementation( + implementation + ) + cell = (self_implementation, cross_implementation) + if cell in latencies: + raise SystemExit( + "Duplicate DiT block benchmark for " + f"{self_implementation} and {cross_implementation}" + ) + latencies[cell] = latency + if self_implementation not in self_configurations: + self_configurations.append(self_implementation) + if cross_implementation not in cross_configurations: + cross_configurations.append(cross_implementation) + + reference_key = ("omnidreams", "omnidreams") + reference = latencies.get(reference_key) + if reference is None: + raise SystemExit("DiT block results have no all-Omnidreams reference") + + matrix = [ + [ + latencies.get((self_implementation, cross_implementation), math.nan) + / reference + for cross_implementation in cross_configurations + # ponytail: Linear sizing assumes current short labels; measure rendered + # text extents if benchmark labels become substantially longer. + ] + for self_implementation in self_configurations + ] + figure, axes = plt.subplots( + figsize=( + max(12.0, len(cross_configurations) * 2.0), + max(8.0, len(self_configurations) * 0.9), + ), + layout="constrained", + ) + image = axes.imshow( + matrix, + aspect="auto", + cmap="RdYlGn_r", + norm=CenteredNorm(vcenter=1.0), + ) + axes.set_xticks( + range(len(cross_configurations)), + labels=[ + _attention_configuration_label(configuration).removeprefix("Optimized ") + for configuration in cross_configurations + ], + rotation=45, + ha="right", + rotation_mode="anchor", + ) + axes.set_yticks( + range(len(self_configurations)), + labels=[ + _attention_configuration_label(configuration).removeprefix("Optimized ") + for configuration in self_configurations + ], + ) + axes.set_xlabel("Cross-attention implementation") + axes.set_ylabel("Self-attention implementation") + axes.set_title(f"Omnidreams DiT block performance\n{subtitle}") + + for self_index, self_implementation in enumerate(self_configurations): + for cross_index, cross_implementation in enumerate(cross_configurations): + latency = latencies.get((self_implementation, cross_implementation)) + if latency is None: + label = "N/A" + else: + label = _bar_label( + latency, + reference, + is_reference=( + self_implementation, + cross_implementation, + ) + == reference_key, + ) + relative_value = matrix[self_index][cross_index] + text_color = "black" + if math.isfinite(relative_value): + red, green, blue, _ = image.cmap(image.norm(np.asarray(relative_value))) + linear = tuple( + channel / 12.92 + if channel <= 0.04045 + else ((channel + 0.055) / 1.055) ** 2.4 + for channel in (red, green, blue) + ) + luminance = 0.2126 * linear[0] + 0.7152 * linear[1] + 0.0722 * linear[2] + text_color = "white" if luminance < 0.179 else "black" + axes.text( + cross_index, + self_index, + label, + ha="center", + va="center", + color=text_color, + ) + + colorbar = figure.colorbar( + image, + ax=axes, + label="Runtime relative to all-Omnidreams reference (×)", + ) + 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, + ) + 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 the three Omnidreams benchmark figures.""" + args = _parse_args(argv) + values, configurations, subtitle = _load_results(args.input) + if args.pipeline_generate_only: + output_path = args.output_dir / "pipeline_generate.png" + _write_bar_figure( + output_path, + "Omnidreams pipeline generate benchmark", + (_PIPELINE_GENERATE_PANEL,), + values, + configurations, + subtitle, + ) + print(f"Wrote {output_path}") + return + + attention_output = args.output_dir / "modules.png" + _write_bar_figure( + attention_output, + "Omnidreams attention module benchmarks", + _ATTENTION_PANELS, + values, + configurations, + subtitle, + ) + print(f"Wrote {attention_output}") + + block_output = args.output_dir / "dit_block.png" + _write_block_figure(block_output, values, subtitle) + print(f"Wrote {block_output}") + + end_to_end_output = args.output_dir / "network_pipeline.png" + _write_bar_figure( + end_to_end_output, + "Omnidreams network and pipeline benchmarks", + _END_TO_END_PANELS, + values, + configurations, + subtitle, + ) + print(f"Wrote {end_to_end_output}") + + +if __name__ == "__main__": + main() diff --git a/scripts/benchmark/omnidreams/run.sh b/scripts/benchmark/omnidreams/run.sh new file mode 100755 index 000000000..be1b7686c --- /dev/null +++ b/scripts/benchmark/omnidreams/run.sh @@ -0,0 +1,16 @@ +#!/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 --package flashdreams-omnidreams python integrations/omnidreams/omnidreams_singleview/tools/sync_thirdparty.py sync + +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 + +uv run python scripts/benchmark/omnidreams/plot.py 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/run_all_plot.sh b/scripts/benchmark/run_all_plot.sh new file mode 100755 index 000000000..ecce02a6c --- /dev/null +++ b/scripts/benchmark/run_all_plot.sh @@ -0,0 +1,16 @@ +#!/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) +cd "$(cd "$script_dir/../.." && pwd)" + +status=0 +for plot_script in "$script_dir"/**/plot*.py; do + uv run python "$plot_script" || status=$? +done + +exit "$status" diff --git a/uv.lock b/uv.lock index d894ab6a8..9e515f96a 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" }, ] @@ -3499,6 +3503,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 +3770,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 +3821,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 +3992,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"