diff --git a/docs/contributing/env_vars.md b/docs/contributing/env_vars.md index 80236691a..339756858 100644 --- a/docs/contributing/env_vars.md +++ b/docs/contributing/env_vars.md @@ -13,7 +13,9 @@ test in the same pull request. them directly with `os.environ.get("NAME")`, and the name must be in the external-variable allowlist (`EXTERNAL_ALLOWLIST` in the contract test). When FastVideo sets such a variable for the other tool, it calls `envs.set_external`, `envs.setdefault_external`, or `envs.unset_external`, and the name must be in - `EXTERNAL_WRITE_ALLOWLIST`. Variables that FastVideo's CI and CI tooling define (for example `TEST_SCOPE` and + `EXTERNAL_WRITE_ALLOWLIST`. Code that needs every variable of a family that another tool owns, such as FA4's + `FA4_*` tuning knobs, calls `envs.external_with_prefixes("FA4_")` with literal prefixes that the allowlist + covers. Variables that FastVideo's CI and CI tooling define (for example `TEST_SCOPE` and `PERF_RUN_SOURCE`) keep their names, and test code under `fastvideo/tests/` reads them directly; they are listed in `CI_ONLY_VARIABLES` in the contract test, together with the file that sets each one. 2. **Read with `envs.NAME.get()`, write with `envs.NAME.set()`, and change a value in tests with diff --git a/fastvideo/attention/utils/fa4_env_cache.py b/fastvideo/attention/utils/fa4_env_cache.py new file mode 100644 index 000000000..218d9be18 --- /dev/null +++ b/fastvideo/attention/utils/fa4_env_cache.py @@ -0,0 +1,71 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Drop FA4's JIT kernel caches when its FA4_*/FA_* tuning environment changes. + +FA4 reads these knobs while tracing a kernel but does not include them in its compile-cache key. +""" + +from __future__ import annotations + +from collections.abc import Iterable, Mapping + +import fastvideo.envs as envs +from fastvideo.logger import init_logger + +logger = init_logger(__name__) + +FA4_ENV_PREFIXES = ("FA4_", "FA_") + +_JIT_CACHE_ATTRS = ("compile_cache", "compile_cache_pre", "compile_cache_post") + +_last_fingerprint: tuple[tuple[str, str], ...] | None = None + + +def fa4_env_fingerprint(environ: Mapping[str, str] | None = None, ) -> tuple[tuple[str, str], ...]: + """Return the sorted ``(name, value)`` pairs of FA4 tuning env vars.""" + source = envs.external_with_prefixes("FA4_", "FA_") if environ is None else environ + pairs = [(name, source[name]) for name in source if name.startswith(FA4_ENV_PREFIXES)] + pairs.sort() + return tuple(pairs) + + +def clear_fa4_jit_caches(kernel_objects: Iterable[object]) -> int: + """Drop every FA4 JIT cache on ``kernel_objects``; return entries dropped.""" + cleared = 0 + for kernel_object in kernel_objects: + for attribute in _JIT_CACHE_ATTRS: + cache = getattr(kernel_object, attribute, None) + if isinstance(cache, dict) and cache: + cleared += len(cache) + cache.clear() + return cleared + + +def reset_fa4_env_fingerprint() -> None: + """Forget the last observed fingerprint (tests / fresh-process semantics).""" + global _last_fingerprint + _last_fingerprint = None + + +def invalidate_fa4_jit_caches_on_env_change( + kernel_objects: Iterable[object], + environ: Mapping[str, str] | None = None, +) -> bool: + """Clear FA4's JIT caches if the FA4 tuning env changed since the last call (always on the first call).""" + global _last_fingerprint + fingerprint = fa4_env_fingerprint(environ) + if fingerprint == _last_fingerprint: + return False + previous, _last_fingerprint = _last_fingerprint, fingerprint + cleared = clear_fa4_jit_caches(kernel_objects) + if previous is not None: + changed = sorted(name for name in {k + for k, _ in previous} | {k + for k, _ in fingerprint} + if dict(previous).get(name) != dict(fingerprint).get(name)) + if cleared: + logger.info( + "FA4 tuning env changed (%s); dropped %d cached JIT kernel(s) " + "so the new setting is not served stale code", ", ".join(changed), cleared) + else: + logger.debug("FA4 tuning env changed (%s); FA4 JIT cache was empty", ", ".join(changed)) + return True diff --git a/fastvideo/attention/utils/flash_attn_cute.py b/fastvideo/attention/utils/flash_attn_cute.py index 4b52da60b..7a7c09fbc 100644 --- a/fastvideo/attention/utils/flash_attn_cute.py +++ b/fastvideo/attention/utils/flash_attn_cute.py @@ -5,6 +5,7 @@ from collections.abc import Callable import torch +from fastvideo.attention.utils.fa4_env_cache import invalidate_fa4_jit_caches_on_env_change from fastvideo.logger import init_logger from fastvideo.platforms import current_platform @@ -57,6 +58,13 @@ except ImportError: # Dynamo unwraps functools caches, so keep this cache inside the opaque helper. _SM90_OR_NEWER_BY_DEVICE: dict[int, bool] = {} +# FA4's JIT cache key ignores its FA4_*/FA_* env knobs; the custom-op impls below run eagerly, so guard there. +_FA4_KERNEL_OBJECTS = (_flash_attn_fwd, _flash_attn_bwd) + + +def _guard_fa4_jit_cache() -> None: + invalidate_fa4_jit_caches_on_env_change(_FA4_KERNEL_OBJECTS) + def _check_dropout(dropout_p: float) -> None: if dropout_p != 0.0: @@ -106,6 +114,7 @@ def _flash_attn_cute_forward( causal: bool, deterministic: bool, ) -> tuple[torch.Tensor, torch.Tensor]: + _guard_fa4_jit_cache() # _flash_attn_fwd returns (out, lse) on its empty-sequence early path but # (out, lse, p, row_max) on the main path at the pinned FA4 cute ref; take # the first two so both arities work. @@ -247,6 +256,7 @@ def _flash_attn_cute_varlen_forward( causal: bool, deterministic: bool, ) -> tuple[torch.Tensor, torch.Tensor]: + _guard_fa4_jit_cache() out, lse = _flash_attn_fwd( q, k, @@ -453,6 +463,7 @@ def _flash_attn_cute_fp4_forward( softmax_scale: float | None, causal: bool, ) -> torch.Tensor: + _guard_fa4_jit_cache() out = _flash_attn_fwd( q, k, diff --git a/fastvideo/envs.py b/fastvideo/envs.py index 8268d95a7..e4ca89460 100644 --- a/fastvideo/envs.py +++ b/fastvideo/envs.py @@ -222,6 +222,15 @@ def unset_external(name: str) -> None: os.environ.pop(name, None) +def external_with_prefixes(*prefixes: str) -> dict[str, str]: + """Return the variables outside the registry whose names start with one of ``prefixes``. + + For tools that read an open-ended family of variables, such as FA4's ``FA4_*`` tuning knobs. The contract test + checks each prefix against its read allowlist. + """ + return {name: value for name, value in os.environ.items() if name.startswith(prefixes)} + + @contextmanager def override_external(name: str, value: str | None) -> Iterator[None]: """Set, or unset when ``value`` is None, a variable outside the registry, and restore it on exit. diff --git a/fastvideo/layers/mxfp8linear.py b/fastvideo/layers/mxfp8linear.py index 0fe017756..6aef7f087 100644 --- a/fastvideo/layers/mxfp8linear.py +++ b/fastvideo/layers/mxfp8linear.py @@ -7,7 +7,16 @@ import torch import torch.nn.functional as F import triton import triton.language as tl -from quack.mx_utils import to_blocked, to_mx + +# Quack moved these helpers to quack.blockscaled.quantize, and some images ship neither. +try: + from quack.mx_utils import to_blocked, to_mx +except ImportError: + try: + from quack.blockscaled.quantize import to_blocked, to_mx + except ImportError: + to_blocked = None + to_mx = None MXFP8_BLOCK_SIZE = 32 _BLOCKS_PER_PROGRAM = 16 @@ -31,6 +40,11 @@ def _quantize_mxfp8_weight_blockwise(matrix: torch.Tensor) -> tuple[torch.Tensor def quantize_mxfp8_weight_blockwise(matrix: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """Prequantize a weight with Quack and return hardware-blocked scales.""" + if to_mx is None or to_blocked is None: + raise ImportError("MXFP8 weight quantization requires Quack's to_mx/to_blocked helpers " + "(quack.mx_utils in older releases, quack.blockscaled.quantize in newer " + "ones). Install a compatible quack to use the MXFP8 weight path; the " + "Triton activation and SwiGLU quantizers do not need it.") _validate_mxfp8_matrix(matrix) return _quantize_mxfp8_weight_blockwise(matrix) @@ -210,6 +224,13 @@ def swiglu_quantize_mxfp8_blockwise(preactivation: torch.Tensor) -> tuple[torch. return quantized, blocked_scales +def _as_e8m0_scales(scales: torch.Tensor) -> torch.Tensor: + # F.scaled_mm rejects flashinfer's uint8 E8M0 scale buffers; they hold the same bytes as float8_e8m0fnu. + if scales.dtype == torch.uint8: + return scales.view(torch.float8_e8m0fnu) + return scales + + def mxfp8_scaled_mm( activation_values: torch.Tensor, activation_scales: torch.Tensor, @@ -221,9 +242,9 @@ def mxfp8_scaled_mm( return F.scaled_mm( mat_a=activation_values, mat_b=weight_values.mT, - scale_a=activation_scales, + scale_a=_as_e8m0_scales(activation_scales), scale_recipe_a=F.ScalingType.BlockWise1x32, - scale_b=weight_scales, + scale_b=_as_e8m0_scales(weight_scales), scale_recipe_b=F.ScalingType.BlockWise1x32, swizzle_a=F.SwizzleType.SWIZZLE_32_4_4, swizzle_b=F.SwizzleType.SWIZZLE_32_4_4, diff --git a/fastvideo/models/dits/minimax_h3.py b/fastvideo/models/dits/minimax_h3.py index 3ce84cdd2..080616fbd 100644 --- a/fastvideo/models/dits/minimax_h3.py +++ b/fastvideo/models/dits/minimax_h3.py @@ -403,16 +403,18 @@ class MiniMaxH3Attention(nn.Module): with STAGES.span("out_proj"): hidden_states, _ = self.to_out(hidden_states) return hidden_states + gate_linear = (self.to_gate_compress if self.to_gate_compress is not None and self._gate_active() else None) with STAGES.span("qkv_proj"): - # All three projections see the same activations. Reuse their FP8 - # quantization when the loaded methods have identical granularity. + # q/k/v and the VSA gate see the same activations; share compatible input quantization. from fastvideo.models.dits.minimax_h3_vsa_fp4 import _shared_input_projections + projections = (self.to_q, self.to_k, self.to_v) + ((gate_linear, ) if gate_linear is not None else ()) if not torch.is_grad_enabled() and not torch.compiler.is_compiling(): - query, key, value = _shared_input_projections((self.to_q, self.to_k, self.to_v), hidden_states) + projected = _shared_input_projections(projections, hidden_states) else: - query, _ = self.to_q(hidden_states) - key, _ = self.to_k(hidden_states) - value, _ = self.to_v(hidden_states) + projected = [projection(hidden_states)[0] for projection in projections] + query, key, value, *gate_out = projected + gate_compress = gate_out[0] if gate_out else None + del projected, gate_out query = query.unflatten(-1, (self.num_attention_heads, self.attention_head_dim)) key = key.unflatten(-1, (self.num_attention_heads, self.attention_head_dim)) value = value.unflatten(-1, (self.num_attention_heads, self.attention_head_dim)) @@ -448,8 +450,7 @@ class MiniMaxH3Attention(nn.Module): # H3 rotates only 96/128 channels, which the generic `freqs_cis` # branch cannot express. Apply it above, then pass no RoPE here. extra_attention_kwargs = {} - if self.to_gate_compress is not None and self._gate_active(): - gate_compress, _ = self.to_gate_compress(hidden_states) + if gate_compress is not None: extra_attention_kwargs["gate_compress"] = gate_compress.unflatten( -1, (self.num_attention_heads, self.attention_head_dim)) with STAGES.span("attention"): diff --git a/fastvideo/models/dits/minimax_h3_vsa_fp4.py b/fastvideo/models/dits/minimax_h3_vsa_fp4.py index 4186cc8bb..921fcf2d1 100644 --- a/fastvideo/models/dits/minimax_h3_vsa_fp4.py +++ b/fastvideo/models/dits/minimax_h3_vsa_fp4.py @@ -186,25 +186,42 @@ def _layout_for(meta: MiniMaxH3VSAMetadata, rotary_emb: tuple[torch.Tensor, torc return layout +def _prequantized_input_key(method: Any, linear: Any) -> tuple[Any, ...] | None: + """Key for linears that can share one pre-quantized input: unit-scale NVFP4, or FP8 by granularity.""" + from fastvideo.layers.quantization.fp8_config import FP8QuantizeMethod + from fastvideo.layers.quantization.nvfp4_config import NVFP4QuantizeMethod + + if type(method) is NVFP4QuantizeMethod: + if method.wants_prequantized_input() and method.uses_unit_activation_scale(linear): + return ("nvfp4", ) + return None + if type(method) is FP8QuantizeMethod: + return ("fp8", method.granularity) + return None + + def _shared_input_projections(linears: tuple[Any, ...], x: torch.Tensor) -> list[torch.Tensor]: """Share compatible FP8 preparation or unit-scale NVFP4 activation quantization. Calibrated or dynamic NVFP4 activation scales retain independent preparation. """ - from fastvideo.layers.quantization.nvfp4_config import NVFP4QuantizeMethod - from fastvideo.layers.quantization.fp8_config import FP8QuantizeMethod - methods = [getattr(linear, "quant_method", None) for linear in linears] - if any(m is None for m in methods): # LoRA wrappers keep quant_method on base_layer. - return [linear(x)[0] for linear in linears] - same_nvfp4 = all(type(m) is NVFP4QuantizeMethod and m.wants_prequantized_input() - and m.uses_unit_activation_scale(linear) - for m, linear in zip(methods, linears, strict=True)) - same_fp8 = all(type(m) is FP8QuantizeMethod and m.granularity == methods[0].granularity for m in methods) - if not (same_nvfp4 or same_fp8): - return [linear(x)[0] for linear in linears] - pre = methods[0].quantize_input(x) - return [m.apply(linear, x, linear.bias, pre_quantized=pre) for m, linear in zip(methods, linears, strict=True)] + keys = [ + _prequantized_input_key(m, linear) if m is not None else None + for m, linear in zip(methods, linears, strict=True) + ] + outputs: list[torch.Tensor | None] = [None] * len(linears) + groups: dict[tuple[Any, ...], list[int]] = {} + for index, key in enumerate(keys): + if key is None: + outputs[index] = linears[index](x)[0] + else: + groups.setdefault(key, []).append(index) + for indices in groups.values(): + pre = methods[indices[0]].quantize_input(x) + for index in indices: + outputs[index] = methods[index].apply(linears[index], x, linears[index].bias, pre_quantized=pre) + return outputs # type: ignore[return-value] @torch.compile(dynamic=True, fullgraph=True) @@ -308,9 +325,11 @@ def vsa_tile_first_attention(attn: Any, hidden_states: torch.Tensor, heads, dim = attn.num_attention_heads, attn.attention_head_dim with STAGES.span("tile_input"): x_tiles = layout.gather_in(hidden_states)[:, :logical] + gate_linear = attn.to_gate_compress if attn.to_gate_compress is not None and attn._gate_active() else None + projections = (attn.to_q, attn.to_k, attn.to_v) + ((gate_linear, ) if gate_linear is not None else ()) with STAGES.span("qkv_proj"): - query, key, value = (t.unflatten(-1, (heads, dim)) - for t in _shared_input_projections((attn.to_q, attn.to_k, attn.to_v), x_tiles)) + query, key, value, *gate_out = _shared_input_projections(projections, x_tiles) + query, key, value = (t.unflatten(-1, (heads, dim)) for t in (query, key, value)) with STAGES.span("qknorm_rope"): cos, sin = layout.cos[:logical], layout.sin[:logical] if use_fused_rope: @@ -320,11 +339,8 @@ def vsa_tile_first_attention(attn: Any, hidden_states: torch.Tensor, else: query = attn._apply_rotary_emb(attn.norm_q(query), (cos, sin)) key = attn._apply_rotary_emb(attn.norm_k(key), (cos, sin)) - gate = None - if attn.to_gate_compress is not None and attn._gate_active(): - with STAGES.span("gate_proj"): - gate, _ = attn.to_gate_compress(x_tiles) - gate = gate.unflatten(-1, (heads, dim)) + gate = gate_out[0].unflatten(-1, (heads, dim)) if gate_out else None + del gate_out capture_root = envs.FASTVIDEO_H3_CAPTURE_QKV.get() if capture_root and attn._layer_idx in (0, 20, 41): from pathlib import Path @@ -536,9 +552,11 @@ def vsa_fp4_attention_sp(attn: Any, hidden_states: torch.Tensor, rotary_emb: tup layout = _SPTileLayout(meta, rank, local_rows) meta._h3_fp4_sp_layout = layout # type: ignore[attr-defined] + gate_linear = attn.to_gate_compress if attn.to_gate_compress is not None and attn._gate_active() else None + projections = (attn.to_q, attn.to_k, attn.to_v) + ((gate_linear, ) if gate_linear is not None else ()) with STAGES.span("qkv_proj_rope"): - query, key, value = (t.unflatten(-1, (heads, dim)) - for t in _shared_input_projections((attn.to_q, attn.to_k, attn.to_v), hidden_states)) + query, key, value, *gate_out = _shared_input_projections(projections, hidden_states) + query, key, value = (t.unflatten(-1, (heads, dim)) for t in (query, key, value)) if use_fused_rope: from fastvideo.models.dits.minimax_h3_fusions import fused_qknorm_rope cos, sin = rotary_emb[0].to(query.dtype), rotary_emb[1].to(query.dtype) @@ -578,7 +596,7 @@ def vsa_fp4_attention_sp(attn: Any, hidden_states: torch.Tensor, rotary_emb: tup with STAGES.span("out_unpack"): out = _unpack_heads_fp8(payload, scale) # [rows, H, D] - if attn.to_gate_compress is not None and attn._gate_active(): + if gate_linear is not None: with STAGES.span("gate_compress"): v_pooled = _pool_tiles(v_t[:, :logical], vbs, layout.tile) out_c = torch.matmul(torch.softmax(scores, dim=-1), v_pooled)[0].to(out.dtype) # [Hs, n_tiles, D] @@ -587,6 +605,5 @@ def vsa_fp4_attention_sp(attn: Any, hidden_states: torch.Tensor, rotary_emb: tup dist.all_gather_into_tensor(gathered, out_c.contiguous(), group=sp_group.device_group) with STAGES.span("gate_apply"): out_c_all = gathered.flatten(0, 1).transpose(0, 1) # [n_tiles, H, D] - gate, _ = attn.to_gate_compress(hidden_states) - out = _apply_gate(out, out_c_all, layout.local_row_tile, gate[0].unflatten(-1, (heads, dim))) + out = _apply_gate(out, out_c_all, layout.local_row_tile, gate_out[0][0].unflatten(-1, (heads, dim))) return out.flatten(1, 2).unsqueeze(0) diff --git a/fastvideo/tests/attention/test_fa4_env_cache.py b/fastvideo/tests/attention/test_fa4_env_cache.py new file mode 100644 index 000000000..aebcce918 --- /dev/null +++ b/fastvideo/tests/attention/test_fa4_env_cache.py @@ -0,0 +1,149 @@ +# SPDX-License-Identifier: Apache-2.0 +"""FA4 JIT-cache invalidation on FA4_*/FA_* tuning-env changes. + +The CPU tests cover the fingerprint + invalidation logic with fake kernel +objects, so they run on any machine. The last test needs a real FA4 install and +a Blackwell/Hopper GPU: it shows that a knob change inside one process reuses +the first arm's kernel without the guard (the bug) and compiles a second kernel +with it (the fix). +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +import fastvideo.envs as envs +from fastvideo.attention.utils import fa4_env_cache + + +@pytest.fixture(autouse=True) +def _fresh_fingerprint(): + fa4_env_cache.reset_fa4_env_fingerprint() + yield + fa4_env_cache.reset_fa4_env_fingerprint() + + +def _fake_kernels(*, fwd_entries: dict | None = None, bwd_entries: dict | None = None): + fwd = SimpleNamespace(compile_cache={} if fwd_entries is None else dict(fwd_entries)) + bwd = SimpleNamespace( + compile_cache={} if bwd_entries is None else dict(bwd_entries), + compile_cache_pre={}, + compile_cache_post={}, + ) + return fwd, bwd + + +def test_fingerprint_selects_only_fa4_tuning_vars() -> None: + environ = { + "FA4_FP8_PV_USE_FUSED_PACK": "1", + "FA_DISABLE_2CTA": "0", + "FASTVIDEO_FA4": "1", # a FastVideo flag, not an FA4 kernel knob + "PATH": "/usr/bin", + } + assert fa4_env_cache.fa4_env_fingerprint(environ) == ( + ("FA4_FP8_PV_USE_FUSED_PACK", "1"), + ("FA_DISABLE_2CTA", "0"), + ) + + +def test_empty_environment_has_empty_fingerprint() -> None: + assert fa4_env_cache.fa4_env_fingerprint({}) == () + + +def test_first_call_drops_preexisting_kernels() -> None: + fwd, bwd = _fake_kernels(fwd_entries={"a": object()}, bwd_entries={"b": object()}) + + cleared = fa4_env_cache.invalidate_fa4_jit_caches_on_env_change((fwd, bwd), {"FA4_X": "1"}) + + assert cleared is True + assert fwd.compile_cache == {} + assert bwd.compile_cache == {} + + +def test_unchanged_environment_keeps_the_compiled_kernel() -> None: + fwd, bwd = _fake_kernels() + environ = {"FA4_X": "1"} + assert fa4_env_cache.invalidate_fa4_jit_caches_on_env_change((fwd, bwd), environ) is True + fwd.compile_cache["a"] = object() + + assert fa4_env_cache.invalidate_fa4_jit_caches_on_env_change((fwd, bwd), environ) is False + assert list(fwd.compile_cache) == ["a"] + + +def test_changed_knob_drops_the_compiled_kernel() -> None: + fwd, bwd = _fake_kernels() + fa4_env_cache.invalidate_fa4_jit_caches_on_env_change((fwd, bwd), {"FA4_X": "1"}) + fwd.compile_cache["a"] = object() + + assert fa4_env_cache.invalidate_fa4_jit_caches_on_env_change((fwd, bwd), {"FA4_X": "0"}) is True + assert fwd.compile_cache == {} + + fwd.compile_cache["a"] = object() + assert fa4_env_cache.invalidate_fa4_jit_caches_on_env_change((fwd, bwd), {"FA4_X": "1"}) is True + assert fwd.compile_cache == {} + + +def test_unrelated_environment_changes_do_not_drop_kernels() -> None: + fwd, bwd = _fake_kernels() + fa4_env_cache.invalidate_fa4_jit_caches_on_env_change((fwd, bwd), {"FA4_X": "1"}) + fwd.compile_cache["a"] = object() + + assert fa4_env_cache.invalidate_fa4_jit_caches_on_env_change((fwd, bwd), { + "FA4_X": "1", + "CUDA_VISIBLE_DEVICES": "3", + }) is False + assert list(fwd.compile_cache) == ["a"] + + +def test_missing_or_non_dict_caches_are_ignored() -> None: + bare = SimpleNamespace() + non_dict = SimpleNamespace(compile_cache=["not", "a", "dict"]) + + assert fa4_env_cache.clear_fa4_jit_caches((bare, non_dict)) == 0 + assert non_dict.compile_cache == ["not", "a", "dict"] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="FA4 needs a CUDA device") +def test_real_fa4_kernel_is_recompiled_for_each_knob_value(monkeypatch) -> None: + """Regression test for the silent stale-kernel reuse. + + Without the guard the second knob value is served the first value's kernel + (one ``cute.compile`` for two settings). With the guard each setting compiles + its own kernel. + """ + interface = pytest.importorskip("flash_attn.cute.interface") + pytest.importorskip("flash_attn.cute.flash_fwd_sm100") + from fastvideo.attention.utils import flash_attn_cute + + torch.manual_seed(0) + q = torch.randn(1, 512, 8, 128, device="cuda", dtype=torch.bfloat16) + k = torch.randn(1, 512, 8, 128, device="cuda", dtype=torch.bfloat16) + v = torch.randn(1, 512, 8, 128, device="cuda", dtype=torch.bfloat16) + + compiles = {"n": 0} + original_compile = interface.cute.compile + + def counting_compile(*args, **kwargs): + compiles["n"] += 1 + return original_compile(*args, **kwargs) + + monkeypatch.setattr(interface.cute, "compile", counting_compile) + + def two_arms() -> int: + interface._flash_attn_fwd.compile_cache.clear() + fa4_env_cache.reset_fa4_env_fingerprint() + compiles["n"] = 0 + for value in ("1", "0"): + with envs.override_external("FA_DISABLE_2CTA", value): + flash_attn_cute.flash_attn_func(q, k, v) + return compiles["n"] + + monkeypatch.setattr(flash_attn_cute, "_guard_fa4_jit_cache", lambda: None) + assert two_arms() == 1 + + monkeypatch.undo() + monkeypatch.setattr(interface.cute, "compile", counting_compile) + assert two_arms() == 2 diff --git a/fastvideo/tests/contract/test_env_policy.py b/fastvideo/tests/contract/test_env_policy.py index 5ee0db55e..6a9230538 100644 --- a/fastvideo/tests/contract/test_env_policy.py +++ b/fastvideo/tests/contract/test_env_policy.py @@ -6,7 +6,8 @@ The test parses every Python file under fastvideo/ (except fastvideo/third_party and the registry fastvideo/envs.py) with ``ast`` and reports code that: - reads the environment directly for a name outside EXTERNAL_ALLOWLIST (test - code under fastvideo/tests/ may also read CI_ONLY_VARIABLES); + code under fastvideo/tests/ may also read CI_ONLY_VARIABLES), or calls + envs.external_with_prefixes with a prefix that EXTERNAL_ALLOWLIST does not cover; - writes the environment directly (os.environ, os.putenv, monkeypatch.setenv), or through the envs.*_external helpers for a name outside EXTERNAL_WRITE_ALLOWLIST; @@ -55,6 +56,8 @@ EXTERNAL_ALLOWLIST = { "NCCL_*": "NCCL communication library.", "TORCH_*": "PyTorch runtime settings.", "OMP_*": "OpenMP threading.", + "FA4_*": "FlashAttention-4 (flash_attn.cute) tuning knobs; fa4_env_cache.py fingerprints them.", + "FA_*": "FlashAttention-4 (flash_attn.cute) tuning knobs; fa4_env_cache.py fingerprints them.", "HF_HUB_*": "huggingface_hub settings.", "RANK": "Set by torchrun and other launchers.", "LOCAL_RANK": "Set by torchrun and other launchers.", @@ -365,6 +368,12 @@ class EnvAccessScanner: elif (isinstance(func, ast.Attribute) and func.attr in EXTERNAL_WRITE_HELPERS and self._is_envs_module(func.value) and not is_external_write_allowed(func.attr, name, self.in_tests)): self.violations.append(("write", name, node.lineno)) + elif (isinstance(func, ast.Attribute) and func.attr == "external_with_prefixes" + and self._is_envs_module(func.value)): + # A prefix read is allowed when an EXTERNAL_ALLOWLIST pattern covers the whole prefix. + for arg in node.args: + prefix = self._literal(arg) + self.violations.append(("read", prefix if prefix == "" else f"{prefix}*", node.lineno)) def _check_registry_use(self, node: ast.Attribute, parent: ast.AST | None, in_function: bool) -> None: # Only a call counts: ``getter = envs.X.get`` neither reads the variable diff --git a/fastvideo/tests/layers/test_mxfp8_optional_quack.py b/fastvideo/tests/layers/test_mxfp8_optional_quack.py new file mode 100644 index 000000000..1e32f9689 --- /dev/null +++ b/fastvideo/tests/layers/test_mxfp8_optional_quack.py @@ -0,0 +1,66 @@ +# SPDX-License-Identifier: Apache-2.0 +"""The in-tree MXFP8 module must import without Quack. + +``fastvideo.layers.mxfp8linear`` used to import ``to_blocked``/``to_mx`` from +``quack.mx_utils`` at module scope. Quack is optional (and newer releases moved +those helpers to ``quack.blockscaled.quantize``), so a missing or too-old Quack +made the whole module -- and everything that imports the MXFP8 layer set -- +unimportable on an otherwise supported Blackwell image. Only weight +prequantization needs Quack. +""" + +from __future__ import annotations + +import builtins +import importlib +import sys + +import pytest +import torch + +_MODULE_NAME = "fastvideo.layers.mxfp8linear" + + +def _import_blocking_quack(monkeypatch: pytest.MonkeyPatch): + """Import the MXFP8 module while every ``quack`` import raises ImportError.""" + real_import = builtins.__import__ + + def blocking_import(name, *args, **kwargs): + if name == "quack" or name.startswith("quack."): + raise ImportError(f"No module named {name!r} (blocked by the test)") + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", blocking_import) + monkeypatch.delitem(sys.modules, _MODULE_NAME, raising=False) + try: + return importlib.import_module(_MODULE_NAME) + finally: + monkeypatch.delitem(sys.modules, _MODULE_NAME, raising=False) + + +def test_mxfp8_module_imports_without_quack(monkeypatch: pytest.MonkeyPatch) -> None: + module = _import_blocking_quack(monkeypatch) + assert module.to_mx is None + assert module.to_blocked is None + assert callable(module.quantize_mxfp8_blockwise) + assert callable(module.swiglu_quantize_mxfp8_blockwise) + assert callable(module.mxfp8_scaled_mm) + + +def test_weight_quantization_without_quack_raises_actionable_error(monkeypatch: pytest.MonkeyPatch) -> None: + module = _import_blocking_quack(monkeypatch) + with pytest.raises(ImportError, match="quack"): + module.quantize_mxfp8_weight_blockwise(torch.zeros(32, 64, dtype=torch.bfloat16)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Quack's MXFP8 helpers compile for CUDA tensors") +def test_weight_quantization_with_quack_matches_blocked_layout() -> None: + pytest.importorskip("quack") + module = importlib.import_module(_MODULE_NAME) + if module.to_mx is None: + pytest.skip("the installed quack has neither mx_utils nor blockscaled.quantize") + values, scales = module.quantize_mxfp8_weight_blockwise(torch.randn(64, 128, dtype=torch.bfloat16, device="cuda")) + assert values.shape == (64, 128) + assert values.dtype == torch.float8_e4m3fn + assert scales.dtype == torch.float8_e8m0fnu + assert scales.numel() % 512 == 0 diff --git a/fastvideo/tests/layers/test_mxfp8_scaled_mm_scale_dtype.py b/fastvideo/tests/layers/test_mxfp8_scaled_mm_scale_dtype.py new file mode 100644 index 000000000..72d688d1e --- /dev/null +++ b/fastvideo/tests/layers/test_mxfp8_scaled_mm_scale_dtype.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +"""MXFP8 GEMMs must accept flashinfer's ``uint8`` E8M0 scale buffers. + +``torch.nn.functional.scaled_mm`` requires ``torch.float8_e8m0fnu`` scales and +rejects the ``uint8`` byte buffers that ``flashinfer.mxfp8_quantize`` returns +with a generic "Invalid scaling configuration" ValueError. FastVideo's own +Triton quantizers already return the e8m0 dtype, so the two paths disagreed on +the same logical scale values. +""" + +from __future__ import annotations + +import pytest +import torch + +from fastvideo.layers.mxfp8linear import ( + _as_e8m0_scales, + mxfp8_scaled_mm, + quantize_mxfp8_blockwise, +) + + +def test_uint8_scales_are_presented_as_e8m0() -> None: + raw = torch.zeros(8, dtype=torch.uint8) + assert _as_e8m0_scales(raw).dtype == torch.float8_e8m0fnu + e8m0 = raw.view(torch.float8_e8m0fnu) + assert _as_e8m0_scales(e8m0) is e8m0 + + +@pytest.mark.skipif(not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 10, + reason="the block-scaled MXFP8 GEMM needs a Blackwell GPU") +def test_mxfp8_scaled_mm_accepts_uint8_and_e8m0_scales() -> None: + torch.manual_seed(0) + x = torch.randn(128, 256, device="cuda", dtype=torch.bfloat16) + w = torch.randn(256, 256, device="cuda", dtype=torch.bfloat16) + activation_values, activation_scales = quantize_mxfp8_blockwise(x) + weight_values, weight_scales = quantize_mxfp8_blockwise(w) + + expected = mxfp8_scaled_mm(activation_values, activation_scales, weight_values, weight_scales, None) + from_uint8 = mxfp8_scaled_mm( + activation_values, + activation_scales.view(torch.uint8), + weight_values, + weight_scales.view(torch.uint8), + None, + ) + torch.testing.assert_close(from_uint8, expected, atol=0, rtol=0) + + reference = x.float() @ w.float().T + relative_error = (expected.float() - reference).norm() / reference.norm() + assert relative_error < 0.05, f"MXFP8 GEMM relative error {relative_error:.4f} is too large" diff --git a/fastvideo/tests/transformers/test_minimax_h3_shared_input_quant.py b/fastvideo/tests/transformers/test_minimax_h3_shared_input_quant.py new file mode 100644 index 000000000..c77d07a0c --- /dev/null +++ b/fastvideo/tests/transformers/test_minimax_h3_shared_input_quant.py @@ -0,0 +1,251 @@ +# SPDX-License-Identifier: Apache-2.0 +"""One activation quantization per compatible H3 attention projection group.""" + +from __future__ import annotations + +from unittest.mock import patch + +import pytest +import torch +import torch.nn as nn + +from fastvideo.models.dits.minimax_h3_vsa_fp4 import _shared_input_projections + + +class _FakeLinear(nn.Module): + """Minimal linear stand-in that records whether its own forward ran.""" + + def __init__(self, method: object) -> None: + super().__init__() + self.quant_method = method + self.bias = None + self.own_forwards = 0 + + def forward(self, x: torch.Tensor): + self.own_forwards += 1 + return (x, None) + + +class _FakeNVFP4Method: + """Version of ``NVFP4QuantizeMethod`` with a countable activation quantization.""" + + def __init__(self, unit_scale: bool = True) -> None: + self.unit_scale = unit_scale + self.quantize_calls = 0 + self.apply_calls = 0 + + def wants_prequantized_input(self) -> bool: + return True + + def uses_unit_activation_scale(self, layer: object) -> bool: + del layer + return self.unit_scale + + def quantize_input(self, x: torch.Tensor): + self.quantize_calls += 1 + return (x + 1, "scale", "global") + + def apply(self, layer: object, x: torch.Tensor, bias: torch.Tensor | None, pre_quantized): + del x, bias + assert pre_quantized is not None + self.apply_calls += 1 + return pre_quantized[0] + + +def test_shared_input_projections_quantize_unit_scale_nvfp4_group_once(monkeypatch: pytest.MonkeyPatch) -> None: + import fastvideo.layers.quantization.nvfp4_config as nvfp4 + + monkeypatch.setattr(nvfp4, "NVFP4QuantizeMethod", _FakeNVFP4Method) + methods = [_FakeNVFP4Method() for _ in range(4)] + linears = tuple(_FakeLinear(method) for method in methods) + x = torch.arange(12, dtype=torch.float32).reshape(3, 4) + + outputs = _shared_input_projections(linears, x) + + assert [method.quantize_calls for method in methods] == [1, 0, 0, 0] + assert [method.apply_calls for method in methods] == [1, 1, 1, 1] + assert [linear.own_forwards for linear in linears] == [0, 0, 0, 0] + for output in outputs: + torch.testing.assert_close(output, x + 1) + + +def test_shared_input_projections_leave_calibrated_nvfp4_separate(monkeypatch: pytest.MonkeyPatch) -> None: + import fastvideo.layers.quantization.nvfp4_config as nvfp4 + + monkeypatch.setattr(nvfp4, "NVFP4QuantizeMethod", _FakeNVFP4Method) + methods = [_FakeNVFP4Method(), _FakeNVFP4Method(), _FakeNVFP4Method(), _FakeNVFP4Method(unit_scale=False)] + linears = tuple(_FakeLinear(method) for method in methods) + x = torch.arange(12, dtype=torch.float32).reshape(3, 4) + + _shared_input_projections(linears, x) + + assert [method.quantize_calls for method in methods] == [1, 0, 0, 0] + assert [method.apply_calls for method in methods] == [1, 1, 1, 0] + assert [linear.own_forwards for linear in linears] == [0, 0, 0, 1] + + +def _install_fp8_buffers(layer, granularity: str, keep_weight: bool = False) -> None: + weight = layer.weight.data.float() + if granularity == "channel": + scale = (weight.abs().amax(dim=1, keepdim=True) / 448).clamp_min(1e-6) + else: + scale = (weight.abs().amax().reshape(1) / 448).clamp_min(1e-6) + layer.register_buffer("_fp8_weight", (weight / scale).to(torch.float8_e4m3fn)) + layer.register_buffer("_fp8_weight_scale", scale) + if not keep_weight: + layer.register_parameter("weight", None) + + +@pytest.mark.parametrize("granularity", ["tensor", "channel"]) +def test_shared_fp8_projections_cover_qkv_and_gate(granularity: str) -> None: + if not torch.cuda.is_available() or torch.cuda.get_device_capability() < (8, 9): + pytest.skip("sm89+ CUDA is required for FP8 GEMM") + from fastvideo.layers.linear import ReplicatedLinear + from fastvideo.layers.quantization.fp8_config import FP8Config, FP8QuantizeMethod + + torch.manual_seed(23) + linears = tuple( + ReplicatedLinear(128, 256, bias=True, quant_config=FP8Config(granularity), prefix=f"block.attn.{name}") + for name in ("to_q", "to_k", "to_v", "to_gate_compress")) + linears[3].quant_method = FP8QuantizeMethod(granularity) + for layer in linears: + layer.to(device="cuda", dtype=torch.bfloat16) + layer.weight.data.normal_(std=0.1) + layer.bias.data.normal_(std=0.1) + _install_fp8_buffers(layer, granularity) + x = torch.randn(1, 272, 128, device="cuda", dtype=torch.bfloat16) + with torch.inference_mode(): + reference = [layer(x)[0] for layer in linears] + with patch.object(FP8QuantizeMethod, "quantize_input", autospec=True, + side_effect=FP8QuantizeMethod.quantize_input) as quant: + actual = _shared_input_projections(linears, x) + assert quant.call_count == 1 + for expected, output in zip(reference, actual, strict=True): + torch.testing.assert_close(output, expected, atol=0, rtol=0) + + +def test_shared_fp8_projections_keep_qkv_sharing_with_dense_gate() -> None: + if not torch.cuda.is_available() or torch.cuda.get_device_capability() < (8, 9): + pytest.skip("sm89+ CUDA is required for FP8 GEMM") + from fastvideo.layers.linear import ReplicatedLinear + from fastvideo.layers.quantization.fp8_config import FP8Config, FP8QuantizeMethod + + torch.manual_seed(29) + linears = tuple( + ReplicatedLinear(128, 256, bias=True, quant_config=FP8Config("tensor"), prefix=f"block.attn.{name}") + for name in ("to_q", "to_k", "to_v")) + linears += (ReplicatedLinear(128, 256, bias=True, prefix="block.attn.to_gate_compress"), ) + for layer in linears[:3]: + layer.to(device="cuda", dtype=torch.bfloat16) + layer.weight.data.normal_(std=0.1) + layer.bias.data.normal_(std=0.1) + _install_fp8_buffers(layer, "tensor") + linears[3].to(device="cuda", dtype=torch.bfloat16) + linears[3].weight.data.normal_(std=0.1) + linears[3].bias.data.normal_(std=0.1) + x = torch.randn(1, 272, 128, device="cuda", dtype=torch.bfloat16) + with torch.inference_mode(): + reference = [layer(x)[0] for layer in linears] + with patch.object(FP8QuantizeMethod, "quantize_input", autospec=True, + side_effect=FP8QuantizeMethod.quantize_input) as quant: + actual = _shared_input_projections(linears, x) + assert quant.call_count == 1 + for expected, output in zip(reference, actual, strict=True): + torch.testing.assert_close(output, expected, atol=0, rtol=0) + + +def test_generic_h3_attention_quantizes_qkv_and_gate_once(env_overrides, distributed_setup) -> None: + """The gate must reuse the q/k/v activation quantization, not prepare it again.""" + if not torch.cuda.is_available() or torch.cuda.get_device_capability() < (8, 9): + pytest.skip("sm89+ CUDA is required for FP8 GEMM") + import fastvideo.layers.quantization.fp8_config as fp8_config + import fastvideo.envs as envs + from fastvideo.attention.backends.video_sparse_attn_h3 import MiniMaxH3VSAMetadataBuilder + from fastvideo.forward_context import set_forward_context + from fastvideo.layers.quantization.fp8_config import FP8Config + from fastvideo.models.dits.minimax_h3 import MiniMaxH3Attention + from fastvideo.platforms import AttentionBackendEnum + + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("VIDEO_SPARSE_ATTN_H3")) + env_overrides.enter_context(envs.override_external("FASTVIDEO_VSA_TRITON", "1")) + env_overrides.enter_context(envs.FASTVIDEO_VSA_PTX_BLACKWELL.override(False)) + env_overrides.enter_context(envs.FASTVIDEO_H3_VSA_FP4.override(False)) + env_overrides.enter_context(envs.FASTVIDEO_H3_VSA_TILE_FIRST.override(False)) + from fastvideo.layers.quantization.fp8_config import FP8QuantizeMethod + + class _GateFP8Config(FP8Config): + """FP8 config that also tags the VSA compression gate (the NVFP4 profile's set).""" + + def get_quant_method(self, layer, prefix): + if prefix.endswith("to_gate_compress"): + from fastvideo.layers.linear import LinearBase + if isinstance(layer, LinearBase): + return FP8QuantizeMethod(self.granularity) + return super().get_quant_method(layer, prefix) + + torch.manual_seed(31) + attn = MiniMaxH3Attention(256, 2, 128, 1e-5, (AttentionBackendEnum.VIDEO_SPARSE_ATTN_H3,), + _GateFP8Config("channel"), "transformer_blocks.0.attn") + attn.to(device="cuda", dtype=torch.bfloat16) + for parameter in attn.parameters(): + parameter.data.normal_(std=0.1) + for layer in (attn.to_q, attn.to_k, attn.to_v, attn.to_out): + _install_fp8_buffers(layer, "channel") + _install_fp8_buffers(attn.to_gate_compress, "channel", keep_weight=True) + meta = MiniMaxH3VSAMetadataBuilder().build(current_timestep=999, patch_size=(1, 1, 1), + VSA_sparsity=0.8, packed_segments=(65, 97, (4, 6, 10)), + device=torch.device("cuda"), tile_size=64) + length = meta.total_seq_length + x = torch.randn(1, length, 256, device="cuda", dtype=torch.bfloat16) + angles = torch.randn(length, 96, device="cuda") + rope = angles.cos(), angles.sin() + with torch.inference_mode(), set_forward_context(current_timestep=0, attn_metadata=meta): + with patch.object(fp8_config, "_quantize_rowwise", side_effect=fp8_config._quantize_rowwise) as quant: + attn(x, rope, length) + assert quant.call_count == 2 + + +@pytest.mark.parametrize("tile_first", [False, True]) +def test_h3_attention_nvfp4_quantizes_qkv_and_gate_once(env_overrides, distributed_setup, tile_first: bool) -> None: + """The NVFP4 ``h3_dit_vsa`` profile prepares q/k/v/gate activations once.""" + if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 10: + pytest.skip("MiniMax-H3 NVFP4 inference requires an NVIDIA Blackwell GPU") + pytest.importorskip("flashinfer") + import fastvideo.envs as envs + import fastvideo.layers.quantization.nvfp4_config as nvfp4 + from fastvideo.attention.backends.video_sparse_attn_h3 import MiniMaxH3VSAMetadataBuilder + from fastvideo.forward_context import set_forward_context + from fastvideo.models.dits.minimax_h3 import MiniMaxH3Attention + from fastvideo.platforms import AttentionBackendEnum + + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("VIDEO_SPARSE_ATTN_H3")) + env_overrides.enter_context(envs.override_external("FASTVIDEO_VSA_TRITON", "1")) + env_overrides.enter_context(envs.FASTVIDEO_VSA_PTX_BLACKWELL.override(False)) + env_overrides.enter_context(envs.FASTVIDEO_H3_VSA_FP4.override(False)) + env_overrides.enter_context(envs.FASTVIDEO_H3_VSA_TILE_FIRST.override(tile_first)) + torch.manual_seed(37) + previous_dtype = torch.get_default_dtype() + torch.set_default_dtype(torch.bfloat16) + try: + attn = MiniMaxH3Attention(256, 2, 128, 1e-5, (AttentionBackendEnum.VIDEO_SPARSE_ATTN_H3,), + nvfp4.NVFP4Config(layer_profile="h3_dit_vsa"), "transformer_blocks.0.attn") + finally: + torch.set_default_dtype(previous_dtype) + attn = attn.cuda().eval() + for parameter in attn.parameters(): + parameter.data.normal_(std=0.1) + nvfp4.convert_model_to_nvfp4(attn) + assert attn.to_q.weight is None + assert attn.to_gate_compress.weight is None + assert attn._gate_active() + meta = MiniMaxH3VSAMetadataBuilder().build(current_timestep=999, patch_size=(1, 1, 1), + VSA_sparsity=0.8, packed_segments=(65, 97, (4, 6, 10)), + device=torch.device("cuda"), tile_size=64) + length = meta.total_seq_length + x = torch.randn(1, length, 256, device="cuda", dtype=torch.bfloat16) + angles = torch.randn(length, 96, device="cuda") + rope = angles.cos(), angles.sin() + with torch.inference_mode(), set_forward_context(current_timestep=0, attn_metadata=meta): + with patch.object(nvfp4, "_nvfp4_quantize", side_effect=nvfp4._nvfp4_quantize) as quant: + attn(x, rope, length) + assert quant.call_count == 2