[bugfix] FA4 JIT cache vs tuning env, MXFP8 import and scale dtype, H3 VSA gate double quantization (#2049)
Co-authored-by: Aryan Kumar <aryan5v@users.noreply.github.com>
This commit is contained in:
Aryan Kumar
co-authored by
Aryan Kumar
Aryan Kumar
parent
0d1244a869
commit
cb6a9a13e1
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -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,
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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"):
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
@@ -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 == "<dynamic>" 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
|
||||
|
||||
@@ -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
|
||||
@@ -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"
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user