From 8db2d71ef217e84bec7dbb6e4ab65c92b144df4f Mon Sep 17 00:00:00 2001 From: posiposi443-ops Date: Sat, 24 Oct 2026 04:40:47 +0900 Subject: [PATCH] [perf] Cache MLX H3 text refinement and speed up attention (#2002) Co-authored-by: Aryan Kumar --- .github/workflows/ci-macos-mlx.yml | 6 + docs/inference/fasth3_omniref_mlx.md | 28 ++ .../inference/basic/mlx_fasth3_omniref.py | 7 +- fastvideo/mlx_runtime/minimax_h3.py | 17 +- fastvideo/mlx_runtime/minimax_h3_pipeline.py | 11 + fastvideo/mlx_runtime/minimax_h3_vsa.py | 32 +- fastvideo/mlx_runtime/minimax_h3_vsa_nax.py | 388 ++++++++++++++++++ fastvideo/mlx_runtime/minimax_h3_vsa_sdpa.py | 97 +++++ .../mlx_runtime/minimax_h3_vsa_target.py | 3 + .../test_minimax_h3_text_refiner_cache.py | 120 ++++++ .../test_minimax_h3_vsa_gather_budget.py | 9 +- .../mlx/test_mlx_minimax_h3_prefix_kv.py | 44 +- .../test_mlx_minimax_h3_ref2va_pipeline.py | 20 +- .../tests/mlx/test_mlx_minimax_h3_vsa.py | 2 +- .../tests/mlx/test_mlx_minimax_h3_vsa_nax.py | 373 +++++++++++++++++ .../tests/mlx/test_mlx_minimax_h3_vsa_sdpa.py | 236 +++++++++++ pyproject.toml | 3 +- 17 files changed, 1382 insertions(+), 14 deletions(-) create mode 100644 fastvideo/mlx_runtime/minimax_h3_vsa_nax.py create mode 100644 fastvideo/mlx_runtime/minimax_h3_vsa_sdpa.py create mode 100644 fastvideo/mlx_runtime/tests/test_minimax_h3_text_refiner_cache.py create mode 100644 fastvideo/tests/mlx/test_mlx_minimax_h3_vsa_nax.py create mode 100644 fastvideo/tests/mlx/test_mlx_minimax_h3_vsa_sdpa.py diff --git a/.github/workflows/ci-macos-mlx.yml b/.github/workflows/ci-macos-mlx.yml index dcb1bee04..3acd1582b 100644 --- a/.github/workflows/ci-macos-mlx.yml +++ b/.github/workflows/ci-macos-mlx.yml @@ -88,7 +88,10 @@ jobs: fastvideo/tests/mlx/test_mlx_affine_dq_gemm.py \ fastvideo/tests/mlx/test_mlx_minimax_h3_parity.py \ fastvideo/tests/mlx/test_mlx_minimax_h3_vsa.py \ + fastvideo/tests/mlx/test_mlx_minimax_h3_vsa_sdpa.py \ + fastvideo/tests/mlx/test_mlx_minimax_h3_vsa_nax.py \ fastvideo/tests/mlx/test_mlx_minimax_h3_vsa_regressions.py \ + fastvideo/tests/mlx/test_mlx_minimax_h3_prefix_kv.py \ fastvideo/tests/mlx/test_mlx_minimax_h3_fast_mode.py \ fastvideo/tests/mlx/test_mlx_minimax_h3_fast_spatial.py \ fastvideo/tests/mlx/test_mlx_minimax_h3_ref_layout.py \ @@ -160,7 +163,10 @@ jobs: fastvideo/tests/mlx/test_mlx_affine_dq_gemm.py \ fastvideo/tests/mlx/test_mlx_minimax_h3_parity.py \ fastvideo/tests/mlx/test_mlx_minimax_h3_vsa.py \ + fastvideo/tests/mlx/test_mlx_minimax_h3_vsa_sdpa.py \ + fastvideo/tests/mlx/test_mlx_minimax_h3_vsa_nax.py \ fastvideo/tests/mlx/test_mlx_minimax_h3_vsa_regressions.py \ + fastvideo/tests/mlx/test_mlx_minimax_h3_prefix_kv.py \ fastvideo/tests/mlx/test_mlx_minimax_h3_fast_mode.py \ fastvideo/tests/mlx/test_mlx_minimax_h3_fast_spatial.py \ fastvideo/tests/mlx/test_mlx_minimax_h3_ref_layout.py \ diff --git a/docs/inference/fasth3_omniref_mlx.md b/docs/inference/fasth3_omniref_mlx.md index eed67f531..29bd25a26 100644 --- a/docs/inference/fasth3_omniref_mlx.md +++ b/docs/inference/fasth3_omniref_mlx.md @@ -73,6 +73,34 @@ Quantized checkpoints pick a linear path per layer. `--quant-gemm` (or The same flag is on the example: `--quant-gemm {auto,dequant,qmm}`. +### Opt-in M5 attention backend + +`--vsa-impl sdpa` evaluates the trained VSA tile mask through fused MLX SDPA, +sharing each attention head's K/V across its query tiles. On M5, the built-in +matrix operations can be faster than the custom SIMD kernel even though they +visit masked keys. The default `auto` selection is unchanged. + +This backend keeps the reference resolution, selected tiles, ragged-row masks, +and compression gate. Accumulation differs from the SIMD backend, so outputs +are not bit-identical. It supports tiles 64/128 with head dimension 128 and +raises if fused SDPA is unavailable. Prefix K/V reuse is not supported; do not +combine it with `--reuse-prefix-kv`. Query tiles are processed in chunks of at +most 128, capped at 256 MiB per boolean-mask chunk with two chunks in flight. + +`--vsa-impl nax` is a separate opt-in backend for Apple M5 GPUs. It uses Metal +TensorOps to visit only the selected key blocks and their valid rows, avoiding +the dense mask and masked-key traversal of the SDPA backend. Dense prefix +queries and the trained compression branch keep their existing implementations. +The selected reference tokens, output resolution and denoising steps are unchanged. + +The NAX backend uses private helper headers from the installed MLX package; +it accepts only the header hashes validated with MLX 0.32.4. Changed headers need +numerical and generation validation before their hashes are accepted. Missing headers, unsupported GPUs +and compilation failures raise errors rather than silently selecting another +backend. Tiles 64/128 and head dimension 128 are supported. Prefix K/V reuse is +not supported. Floating-point accumulation differs from both SIMD and SDPA, +so the generated clip can differ even with the same seed. `auto` remains unchanged. + ## Generate ```bash diff --git a/examples/inference/basic/mlx_fasth3_omniref.py b/examples/inference/basic/mlx_fasth3_omniref.py index ddda8a17a..057c43821 100644 --- a/examples/inference/basic/mlx_fasth3_omniref.py +++ b/examples/inference/basic/mlx_fasth3_omniref.py @@ -114,9 +114,10 @@ def build_parser() -> argparse.ArgumentParser: parser.add_argument("--vsa-ref-keep-rate", type=float, default=None, help="fraction of each reference video's tiles a query keeps (p2_multi_region); " "default is the contract's trained value") - parser.add_argument("--vsa-impl", choices=("auto", "reference", "simd"), default="auto", + parser.add_argument("--vsa-impl", choices=("auto", "reference", "simd", "sdpa", "nax"), default="auto", help="sparse attention backend; auto uses the SIMD Metal kernel for 128-token tiles and " - "gather+SDPA otherwise, reference forces gather+SDPA") + "gather+SDPA otherwise, reference forces gather+SDPA; sdpa shares K/V with a tile mask " + "and nax uses M5 TensorOps over selected blocks (both opt-in, incompatible with --reuse-prefix-kv)") parser.add_argument("--reuse-prefix-kv", action="store_true", help="opt-in, lossy: record the text/reference rows' per-block K/V on step 1 and run later " "steps on the target audio/video rows only (faster, slightly different output)") @@ -145,6 +146,8 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace: parser.error("Ref2VA needs at least one image or video reference") if args.model_root is None and (args.export is None or args.base_model_path is None): parser.error("give --model-root, or --export together with --base-model-path") + if args.vsa_impl in ("sdpa", "nax") and args.reuse_prefix_kv: + parser.error(f"--vsa-impl {args.vsa_impl} does not support --reuse-prefix-kv") return args diff --git a/fastvideo/mlx_runtime/minimax_h3.py b/fastvideo/mlx_runtime/minimax_h3.py index 01d94ab42..bf4d66866 100644 --- a/fastvideo/mlx_runtime/minimax_h3.py +++ b/fastvideo/mlx_runtime/minimax_h3.py @@ -1603,6 +1603,7 @@ class MLXMiniMaxH3DiT: pdd_block: tuple[int, int] | None = None, pdd_integration_weights: dict[str, np.ndarray] | None = None, record_prefix_kv: Any | None = None, + precomputed_text_embeds: Any | None = None, ): """Denoise-step forward served entirely from the AdaLN cache. @@ -1616,6 +1617,11 @@ class MLXMiniMaxH3DiT: ``step_timesteps``. With a PDD checkpoint pass the step's ``pdd_block`` and ``pdd_integration_weights`` so the widened output heads are fused into the block's weighted mean. + + ``precomputed_text_embeds`` may hold ``refine_text(text_rows)`` for + this request. The refiner is independent of the denoising timestep; + callers own its lifetime and must recompute it for changed text or + weights. No embeddings are retained on the model. """ import mlx.core as mx @@ -1640,7 +1646,16 @@ class MLXMiniMaxH3DiT: self.weights["audio_proj_in.weight"], self.weights["audio_proj_in.bias"], ) - text_embeds = self.refine_text(text_rows) + if precomputed_text_embeds is None: + text_embeds = self.refine_text(text_rows) + else: + expected_shape = (len(layout.text_indices), self.hidden_size) + if precomputed_text_embeds.shape != expected_shape: + raise ValueError(f"Precomputed text embeddings have shape {precomputed_text_embeds.shape}; " + f"expected {expected_shape}.") + if precomputed_text_embeds.dtype not in (mx.float32, mx.float16, mx.bfloat16): + raise ValueError("Precomputed text embeddings must have a floating-point dtype.") + text_embeds = precomputed_text_embeds packed = mx.zeros((sequence_length, self.hidden_size), dtype=text_embeds.dtype) packed[mx.array(layout.text_indices)] = text_embeds diff --git a/fastvideo/mlx_runtime/minimax_h3_pipeline.py b/fastvideo/mlx_runtime/minimax_h3_pipeline.py index 26659a91f..ad3b54674 100644 --- a/fastvideo/mlx_runtime/minimax_h3_pipeline.py +++ b/fastvideo/mlx_runtime/minimax_h3_pipeline.py @@ -937,6 +937,10 @@ class MiniMaxH3MLXPipeline: elif layout.num_condition_audio_rows: raise ValueError("A reference-prefixed layout needs its encoded audio condition rows.") text = mx.array(text_rows.astype(np.float32)) + # The token refiner depends only on this request's conditioning, not + # the changing latents or timesteps. A local value also prevents a + # resident DiT from reusing embeddings from an earlier request. + text_embeds = None dit_forward_s = 0.0 prefix_kv = None if reuse_prefix_kv: @@ -998,6 +1002,9 @@ class MiniMaxH3MLXPipeline: if record_step: pdd_forward_kwargs["record_prefix_kv"] = prefix_kv recorded_steps.append(step_index) + if text_embeds is None: + text_embeds = dit.refine_text(text) + mx.eval(text_embeds) video_velocity, audio_velocity = dit.forward_with_cache( x_v, x_a, @@ -1006,6 +1013,7 @@ class MiniMaxH3MLXPipeline: step_timesteps=unique, row_timestep_inverse=inverse, step_index=step_index, + precomputed_text_embeds=text_embeds, **pdd_forward_kwargs, ) video_velocity = video_velocity[layout.num_condition_video_rows:] @@ -1326,6 +1334,9 @@ class MiniMaxH3MLXPipeline: timings: dict[str, float] = {} peaks: dict[str, float] = {} + if vsa_impl in ("sdpa", "nax") and reuse_prefix_kv: + raise ValueError( + f"{vsa_impl.upper()} VSA does not support prefix K/V reuse; disable reuse or select another backend") if fast_sharpen < 0: raise ValueError(f"fast_sharpen must be non-negative, got {fast_sharpen}.") if inter_step_cooldown_s < 0: diff --git a/fastvideo/mlx_runtime/minimax_h3_vsa.py b/fastvideo/mlx_runtime/minimax_h3_vsa.py index a0d15e3e1..589821da2 100644 --- a/fastvideo/mlx_runtime/minimax_h3_vsa.py +++ b/fastvideo/mlx_runtime/minimax_h3_vsa.py @@ -23,6 +23,10 @@ Execution backends: - **simd** — flash-style SIMD-group 8x8 matrix operations over the reference tile map for tile sizes 64 and 128 and head dimension 128. Unsupported shapes and kernel failures fall back to the reference backend. +- **sdpa** — opt-in fused SDPA with the same tile mask and shared K/V. Intended + for M5; visits masked keys but can use faster matrix operations. ``auto`` is unchanged. +- **nax** — opt-in M5 TensorOps kernel that visits only selected blocks and valid + key rows. Uses the installed MLX NAX helpers; ``auto`` is unchanged. Dense fused SDPA remains the default when VSA is disabled, the geometry is unsupported, or a dense-only checkpoint is loaded. @@ -50,7 +54,7 @@ VSA_H3_TILE_SHAPES: dict[int, tuple[int, int, int]] = { VSA_REF_POLICIES = ("p2_multi_region", ) VSA_GATE_KEY_SUFFIX = "attn.to_gate_compress.weight" VSA_PREFIX_MODES = ("exempt", "compete") -PUBLIC_VSA_IMPLS = ("auto", "reference", "simd") +PUBLIC_VSA_IMPLS = ("auto", "reference", "simd", "sdpa", "nax") VSA_IMPLS = PUBLIC_VSA_IMPLS # Tile sizes where ``auto`` picks the SIMD kernel (trained Omni tiles); other tiles stay opt-in. _AUTO_SIMD_TILES: tuple[int, ...] = (128, ) @@ -58,7 +62,7 @@ _REFERENCE_FULL_MASK_TILE_LIMIT = 24 _REFERENCE_FULL_MASK_MAX_ELEMENTS = 256 * 1024**2 PrefixMode = Literal["exempt", "compete"] -VSAImpl = Literal["auto", "reference", "simd"] +VSAImpl = Literal["auto", "reference", "simd", "sdpa", "nax"] def validate_ref_policy(ref_policy: str | None, ref_keep_rate: float | None) -> None: @@ -873,6 +877,20 @@ def _gate_compress_output(scores, v_tiled, gate_tiled, geometry: MiniMaxH3VSAGeo def resolve_impl(requested: VSAImpl, head_dim: int, tile_elems: int) -> str: + if requested == "nax": + from fastvideo.mlx_runtime.minimax_h3_vsa_nax import NAX_HEAD_DIM, NAX_TILE_SIZES + + if tile_elems not in NAX_TILE_SIZES or head_dim != NAX_HEAD_DIM: + raise ValueError("NAX VSA requires tile 64 or 128 and head dimension 128 " + f"(got tile={tile_elems}, head_dim={head_dim})") + return "nax" + if requested == "sdpa": + from fastvideo.mlx_runtime.minimax_h3_vsa_sdpa import SDPA_HEAD_DIM, SDPA_TILE_SIZES + + if tile_elems not in SDPA_TILE_SIZES or head_dim != SDPA_HEAD_DIM: + raise ValueError("SDPA VSA requires tile 64 or 128 and head dimension 128 " + f"(got tile={tile_elems}, head_dim={head_dim})") + return "sdpa" if requested == "reference" or (requested == "auto" and tile_elems not in _AUTO_SIMD_TILES): return "reference" @@ -987,6 +1005,16 @@ def h3_vsa_attention( if stats is not None: stats.dense_fallback_reason = f"simd kernel failed: {error}" video_tiled = _reference_gather_sdpa(q_tiled, k_tiled, v_tiled, block_idx, geometry, scale) + elif chosen == "nax": + from fastvideo.mlx_runtime.minimax_h3_vsa_nax import nax_block_sparse + + video_tiled = nax_block_sparse(q_tiled, k_tiled, v_tiled, block_idx, block_num, geometry, + scale)[n_prefix_pad:] + elif chosen == "sdpa": + from fastvideo.mlx_runtime.minimax_h3_vsa_sdpa import sdpa_block_sparse + + video_tiled = sdpa_block_sparse(q_tiled, k_tiled, v_tiled, block_idx, block_num, geometry, + scale)[n_prefix_pad:] elif _reference_full_mask_fits(geometry, heads): scores_np = np.array(scores, dtype=np.float32) mask = build_block_mask( diff --git a/fastvideo/mlx_runtime/minimax_h3_vsa_nax.py b/fastvideo/mlx_runtime/minimax_h3_vsa_nax.py new file mode 100644 index 000000000..e60e618a5 --- /dev/null +++ b/fastvideo/mlx_runtime/minimax_h3_vsa_nax.py @@ -0,0 +1,388 @@ +# SPDX-License-Identifier: Apache-2.0 AND MIT +# mypy: disable-error-code=no-untyped-call +# +# This file contains a kernel adapted from MLX. Its MIT notice follows. +# MIT License +# +# Copyright © 2023 Apple Inc. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +"""Opt-in M5 block-sparse TensorOps attention using installed MLX NAX helpers. + +Only selected blocks and their valid rows are visited, without K/V duplication +or a dense mask. Prefix and padded query outputs remain zero. Online softmax +and matrix accumulation can differ numerically from SIMD and fused SDPA. + +This backend depends on private headers from the installed MLX package, tested +with MLX 0.32.4 and the Metal 4 SDK. Only the validated header hashes are accepted. +Missing headers, incompatible APIs and +unsupported hardware raise explicitly; no fallback changes the implementation. +Importing this module performs no MLX operations. Header copyrights are retained +and get_nax_provenance() provides hashes for reproducible benchmark receipts. +""" +from __future__ import annotations + +import hashlib +import math +import re +from functools import lru_cache +from pathlib import Path + +NAX_TILE_SIZES = (64, 128) +NAX_HEAD_DIM = 128 +DEFAULT_QUERY_ROWS = 32 +_NAX_HEADER = "mlx/backend/metal/kernels/steel/attn/nax.h" +# Private fragment layouts can change without renaming their APIs. Accept only +# the headers used by the numerical and full-generation validation on MLX 0.32.4. +_VALIDATED_HEADER_SHA256 = { + _NAX_HEADER: + "fb10bcae44095e9fed5c81213d03527dd5fe5ac11be11ec7784bd357c6f317ec", + "mlx/backend/metal/kernels/steel/defines.h": + "b03cea6a7d5cfe814e6838d8536e2aacdcafd2744cb408e954fda336ce21759a", + "mlx/backend/metal/kernels/steel/utils/integral_constant.h": + "886eccaac4f8f00eab5a10551fdb5315df7ab7462aba13d69bb95765c1a4735e", + "mlx/backend/metal/kernels/steel/utils/type_traits.h": + "4766c5a3809d3cf7740e63acc5703ef76b6b4e85d8baf16602c9d517a2fd426c", +} + +_OPERATORS = r""" +namespace h3_sparse_nax { +struct MaxOp { + template + METAL_FUNC static constexpr V apply(V x, V y) { return metal::max(x, y); } +}; +struct SumOp { + template + METAL_FUNC static constexpr V apply(V x, V y) { return x + y; } +}; +struct MulOp { + template + METAL_FUNC static constexpr V apply(V x, V y) { return x * y; } +}; +struct ExpSubOp { + template + METAL_FUNC static constexpr V apply(V x, V y) { return metal::fast::exp2(x - y); } +}; +} +""" + +# Adapted from MLX steel_attention_nax.h (Copyright 2024-26 Apple Inc.). +# The fragment math/online softmax follows that implementation; this launch +# addresses the existing H3 tiled layout and walks selected key blocks directly. +_SOURCE = r""" + using namespace mlx::steel; + using namespace h3_sparse_nax; + constexpr int D = 128; + constexpr int BK = 32; + constexpr int TK = BK / 16; + constexpr int TD = D / 16; + constexpr float NEG = -3.402823466e38f; + + const int qsub = int(threadgroup_position_in_grid.x); + const int query_tile = int(threadgroup_position_in_grid.y); + const int head = int(threadgroup_position_in_grid.z); + const int sid = int(simdgroup_index_in_threadgroup); + const int k_max = meta_i[0]; + const int n_prefix = meta_i[1]; + const int n_video = meta_i[2]; + const int heads = meta_i[3]; + const int qt = n_prefix + query_tile; + const int q_valid = vbs[qt]; + const int row_base = qsub * BQ; + const int nsel = block_num[head * n_video + query_tile]; + + // Uniform across the threadgroup: later barriers cannot be skipped by only + // some SIMD groups. Inactive SIMD groups use guarded loads and stores. + if (row_base >= q_valid || nsel <= 0) { return; } + + const int row_stride = heads * D; + const int qrow = row_base + sid * 16; + const short q_remaining = short(q_valid - qrow); + const device T* Q = q + (long(qt * TILE + qrow) * heads + head) * D; + device T* O = out + (long(qt * TILE + qrow) * heads + head) * D; + const float scale2 = meta_f[0] * 1.44269504088896341f; + + using otile_t = NAXTile; + using stile_t = NAXTile; + otile_t Otile; + Otile.clear(); + constexpr short kRowsPT = otile_t::kRowsPerThread; + metal::vec max_score; + metal::vec sum_score{0}; + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kRowsPT; ++i) { max_score[i] = NEG; } + const short sn = otile_t::NAXFrag_t::get_coord().x; + + for (int selected = 0; selected < nsel; ++selected) { + const int kt = block_idx[(head * n_video + query_tile) * k_max + selected]; + const int k_valid = vbs[kt]; + const long kv_offset = (long(kt) * TILE * heads + head) * D; + for (int krow = 0; krow < k_valid; krow += BK) { + const short k_remaining = short(metal::min(BK, k_valid - krow)); + const device T* K = k + kv_offset + long(krow) * row_stride; + const device T* V = v + kv_offset + long(krow) * row_stride; + stile_t Stile; + Stile.clear(); + + // QK: retain MLX's four-way partial unroll for the 128-wide head. + STEEL_PRAGMA_UNROLL + for (short ik = 0; ik < TK; ik += 2) { +#pragma clang loop unroll_count(4) + for (short id = 0; id < TD; ++id) { + NAXTile Qtile; + NAXTile Ktile; + if (q_remaining < 16) { + Qtile.load_rows(Q + id * 16, row_stride, q_remaining); + } else { + Qtile.load(Q + id * 16, row_stride); + } + if (k_remaining < BK) { + Ktile.load_rows(K + ik * 16 * row_stride + id * 16, + row_stride, k_remaining - ik * 16); + } else { + Ktile.load(K + ik * 16 * row_stride + id * 16, row_stride); + } + stile_t::NAXFrag_t::mma( + Stile.frag_at(0, ik), Stile.frag_at(0, ik + 1), + Qtile.frag_at(0, 0), metal::false_type{}, + Ktile.frag_at(0, 0), Ktile.frag_at(1, 0), metal::true_type{}); + } + } + STEEL_PRAGMA_UNROLL + for (short i = 0; i < stile_t::kElemsPerTile; ++i) { + Stile.elems()[i] *= scale2; + } + if (k_remaining < BK) { + STEEL_PRAGMA_UNROLL + for (short ik = 0; ik < TK; ++ik) { + thread auto& frag = Stile.frag_at(0, ik); + STEEL_PRAGMA_UNROLL + for (short i = 0; i < stile_t::kFragThrRows; ++i) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < stile_t::kFragThrCols; ++j) { + if (ik * 16 + sn + j >= k_remaining) { + frag[i * stile_t::kFragThrCols + j] = NEG; + } + } + } + } + } + + metal::vec new_max = max_score; + metal::vec factor; + Stile.template row_reduce(new_max); + Stile.template row_bin_op(new_max); + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kRowsPT; ++i) { + factor[i] = metal::fast::exp2(max_score[i] - new_max[i]); + max_score[i] = new_max[i]; + sum_score[i] *= factor[i]; + } + Stile.template row_reduce(sum_score); + Otile.template row_bin_op(factor); + simdgroup_barrier(mem_flags::mem_none); + + // PV retains float32 probabilities, matching MLX's NAX SDPA kernel. + STEEL_PRAGMA_UNROLL + for (short id = 0; id < TD; id += 2) { + if (id == 4) { threadgroup_barrier(mem_flags::mem_none); } + STEEL_PRAGMA_UNROLL + for (short ik = 0; ik < TK; ++ik) { + NAXTile Vtile; + if (k_remaining < BK) { + Vtile.load_rows(V + ik * 16 * row_stride + id * 16, + row_stride, k_remaining - ik * 16); + } else { + Vtile.load(V + ik * 16 * row_stride + id * 16, row_stride); + } + otile_t::NAXFrag_t::mma( + Otile.frag_at(0, id), Otile.frag_at(0, id + 1), + Stile.frag_at(0, ik), metal::false_type{}, + Vtile.frag_at(0, 0), Vtile.frag_at(0, 1), metal::false_type{}); + } + } + } + } + threadgroup_barrier(mem_flags::mem_none); + metal::vec rcp; + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kRowsPT; ++i) { + rcp[i] = sum_score[i] > 0.f ? 1.f / sum_score[i] : 0.f; + } + Otile.template row_bin_op(rcp); + if (q_remaining > 0) { + if (q_remaining < 16) { Otile.store_rows(O, row_stride, q_remaining); } + else { Otile.store(O, row_stride); } + } +""" + + +def _mlx_include_root() -> Path: + """Resolve headers beside the imported package, including namespace installs.""" + import mlx + + for package_path in mlx.__path__: + include_root = Path(package_path) / "include" + if (include_root / _NAX_HEADER).is_file(): + return include_root + raise RuntimeError("NAX VSA requires installed MLX NAX headers (tested with MLX 0.32.4); " + "the mlx package has no include/mlx/backend/metal/kernels/steel/attn/nax.h") + + +@lru_cache(maxsize=1) +def _header(): + """Inline quoted MLX includes; system Metal/MPP includes stay intact.""" + include_root = _mlx_include_root() + visited = {} + + def expand(relative): + if relative in visited: + return f"// Already included: {relative}\n" + path = include_root / relative + try: + original_bytes = path.read_bytes() + except OSError: + raise RuntimeError(f"NAX VSA cannot read required installed MLX header: {relative}") from None + digest = hashlib.sha256(original_bytes).hexdigest() + if digest != _VALIDATED_HEADER_SHA256.get(relative): + raise RuntimeError(f"Unvalidated MLX private header: {relative}; NAX VSA requires the headers " + "validated with MLX 0.32.4. Revalidate this backend before accepting new headers.") + original = original_bytes.decode("utf-8") + visited[relative] = digest + content = re.sub(r"^\s*#pragma once\s*$", "", original, flags=re.MULTILINE) + content = re.sub(r'^\s*#include "(mlx/[^"\n]+)"\s*$', + lambda match: expand(match.group(1)), + content, + flags=re.MULTILINE) + if relative == _NAX_HEADER: + old = "transpose_b,\n true,\n mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate" + expected_api = ("struct BaseNAXFrag", "struct NAXTile", "get_coord()", "row_reduce(", "row_bin_op(", + "load_rows(", "store_rows(") + if content.count(old) != 2 or any(marker not in content for marker in expected_api): + raise RuntimeError("Installed MLX NAX private API changed; this backend was tested with MLX 0.32.4 " + "and needs compatibility review") + return f"\n// Begin installed MLX header: {relative}\n{content}\n// End: {relative}\n" + + header = "#ifndef METAL_FUNC\n#define METAL_FUNC inline __attribute__((always_inline))\n#endif\n" + header += expand(_NAX_HEADER) + header += _OPERATORS + return header, visited + + +def get_nax_provenance(): + """CPU-only metadata, containing no machine-specific absolute paths.""" + header, dependencies = _header() + return { + "kernel": "h3-true-sparse-nax-v1", + "relaxed_precision": True, + "math_mode": "safe", + "key_rows": 32, + "source_sha256": hashlib.sha256(_SOURCE.encode()).hexdigest(), + "expanded_header_sha256": hashlib.sha256(header.encode()).hexdigest(), + "installed_header_sha256": dependencies + } + + +def _check_nax_device() -> None: + """The TensorOps backend is deliberately restricted to validated M5 GPUs.""" + import mlx.core as mx + + if not mx.metal.is_available(): + raise RuntimeError("NAX VSA requires an Apple M5 GPU and Metal 4") + device_name = str(mx.device_info().get("device_name", "")) + if re.match(r"^Apple M5(?:\s|$)", device_name) is None: + raise RuntimeError(f"NAX VSA supports Apple M5 GPUs only, got {device_name or 'unknown GPU'}") + + +@lru_cache(maxsize=1) +def _nax_kernel(): + import mlx.core as mx + + _check_nax_device() + header, _ = _header() + if not hasattr(mx.fast, "metal_kernel"): + raise RuntimeError("NAX VSA requires mx.fast.metal_kernel with safe math compile_options") + try: + return mx.fast.metal_kernel(name="h3_vsa_sparse_nax", + input_names=["q", "k", "v", "block_idx", "block_num", "vbs", "meta_i", "meta_f"], + output_names=["out"], + source=_SOURCE, + header=header, + compile_options={"math_mode": "safe"}) + except (TypeError, ValueError, RuntimeError) as error: + raise RuntimeError("NAX VSA kernel construction failed; this backend requires compatible MLX NAX headers " + "and Metal 4 TensorOps support") from error + + +def nax_block_sparse(q_tiled, + k_tiled, + v_tiled, + block_idx, + block_num, + geometry, + scale: float, + *, + query_rows: int = DEFAULT_QUERY_ROWS): + """Same [S_pad,H,128] contract as the existing H3 sparse implementations. + + query_rows controls SIMD groups per threadgroup. + Active indices must be valid and unique, as emitted by the H3 router. Values + in inactive index slots are never read. Prefix/padding outputs are zero. + """ + import mlx.core as mx + + if query_rows not in (16, 32, 64, 128): + raise ValueError("query_rows must be 16, 32, 64 or 128") + if q_tiled.ndim != 3 or q_tiled.shape != k_tiled.shape or q_tiled.shape != v_tiled.shape: + raise ValueError("NAX sparse attention requires matching [S_pad,H,D] Q/K/V") + padded, heads, dim = q_tiled.shape + tile = geometry.tile_elems + if tile not in NAX_TILE_SIZES or dim != NAX_HEAD_DIM or query_rows > tile: + raise ValueError("NAX sparse attention requires tile 64/128, head dim 128, query_rows <= tile") + if padded != geometry.padded_length or heads < 1: + raise ValueError("Q/K/V shape does not match geometry") + if q_tiled.dtype not in (mx.float32, mx.float16, mx.bfloat16): + raise ValueError("NAX sparse attention supports float32, float16 or bfloat16") + if q_tiled.dtype != k_tiled.dtype or q_tiled.dtype != v_tiled.dtype: + raise ValueError("Q/K/V dtypes must match") + if (block_idx.ndim != 3 or block_idx.shape[:2] != (heads, geometry.num_video_tiles) + or block_num.shape != (heads, geometry.num_video_tiles)): + raise ValueError("Routing shapes do not match head/tile geometry") + if not math.isfinite(scale): + raise ValueError("scale must be finite") + if geometry.num_video_tiles == 0: + return mx.zeros_like(q_tiled) + threads = (query_rows // 16) * 32 + result, = _nax_kernel()(inputs=[ + mx.contiguous(q_tiled), + mx.contiguous(k_tiled), + mx.contiguous(v_tiled), + mx.contiguous(block_idx.astype(mx.int32)), + mx.contiguous(block_num.astype(mx.int32)), + mx.array(geometry.variable_block_sizes, dtype=mx.int32), + mx.array([block_idx.shape[-1], geometry.num_prefix_tiles, geometry.num_video_tiles, heads], dtype=mx.int32), + mx.array([scale], dtype=mx.float32) + ], + template=[("T", q_tiled.dtype), ("TILE", tile), ("BQ", query_rows)], + grid=(threads * (tile // query_rows), geometry.num_video_tiles, heads), + threadgroup=(threads, 1, 1), + output_shapes=[q_tiled.shape], + output_dtypes=[q_tiled.dtype], + init_value=0) + return result diff --git a/fastvideo/mlx_runtime/minimax_h3_vsa_sdpa.py b/fastvideo/mlx_runtime/minimax_h3_vsa_sdpa.py new file mode 100644 index 000000000..a6e342cbb --- /dev/null +++ b/fastvideo/mlx_runtime/minimax_h3_vsa_sdpa.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: Apache-2.0 +# mypy: disable-error-code=no-untyped-call +"""VSA tile masks through fused SDPA, sharing K/V with grouped-query attention. + +This opt-in backend preserves selected blocks, ragged-row masks and the trained +compression branch. Unlike the SIMD kernel, fused SDPA visits masked key blocks; +on M5 its faster matrix operations can outweigh that additional work. Floating- +point accumulation differs, so this is not a bit-identical implementation. +""" + +from __future__ import annotations + +SDPA_TILE_SIZES = (64, 128) +SDPA_HEAD_DIM = 128 +DEFAULT_QUERY_CHUNK = 128 +_MASK_CHUNK_MAX_BYTES = 256 * 1024**2 + + +def sdpa_block_sparse(q_tiled, + k_tiled, + v_tiled, + block_idx, + block_num, + geometry, + scale: float, + *, + chunk_size: int = DEFAULT_QUERY_CHUNK): + """Return tiled [S_pad,H,D], with zero prefix rows and zero padded queries. + + Query heads are ordered (original head, query tile), so GQA maps each group + of query tiles to its original K/V head. A tile's key mask broadcasts across + its query rows. At most two mask chunks are in flight. Unsupported fused + operations raise instead of silently allocating unfused attention scores. + """ + import mlx.core as mx + + if chunk_size < 1: + raise ValueError("SDPA VSA chunk_size must be positive") + if q_tiled.ndim != 3 or k_tiled.shape != q_tiled.shape or v_tiled.shape != q_tiled.shape: + raise ValueError("SDPA VSA requires matching [S_pad,H,D] Q/K/V shapes") + padded, heads, dim = q_tiled.shape + tile = geometry.tile_elems + if tile not in SDPA_TILE_SIZES or dim != SDPA_HEAD_DIM: + raise ValueError(f"SDPA VSA requires tile 64 or 128 and head dimension 128, got tile={tile}, dim={dim}") + if padded != geometry.padded_length or heads < 1: + raise ValueError("SDPA VSA Q/K/V shape does not match the tile geometry") + if q_tiled.dtype != k_tiled.dtype or q_tiled.dtype != v_tiled.dtype: + raise ValueError("SDPA VSA requires matching Q/K/V dtypes") + n_tiles = geometry.num_tiles + n_video = geometry.num_video_tiles + n_prefix = geometry.num_prefix_tiles + if (block_idx.ndim != 3 or block_idx.shape[:2] != (heads, n_video) or block_num.shape != (heads, n_video)): + raise ValueError("SDPA VSA routing shape does not match the query head/tile geometry") + + query_chunk = min(chunk_size, max(1, _MASK_CHUNK_MAX_BYTES // max(heads * padded, 1))) + q_video = q_tiled.reshape(n_tiles, tile, heads, dim).transpose(2, 0, 1, 3)[:, n_prefix:] + # Fused full attention accepts these strided views: the D axis is contiguous. + k_shared = k_tiled.transpose(1, 0, 2)[None] + v_shared = v_tiled.transpose(1, 0, 2)[None] + sizes = mx.array(geometry.variable_block_sizes, dtype=mx.int32) + row = mx.arange(tile, dtype=mx.int32) + key_valid = row[None, :] < sizes[:, None] + head_index = mx.arange(heads, dtype=mx.int32)[:, None, None] + selected_slot = mx.arange(block_idx.shape[-1], dtype=mx.int32)[None, None, :] + outputs = [] + for start in range(0, n_video, query_chunk): + end = min(start + query_chunk, n_video) + count = end - start + indices = block_idx[:, start:end].astype(mx.int32) + counts = block_num[:, start:end] + # Send inactive slots to a sentinel column that is discarded. Every + # scatter writes True, so duplicate indices cannot erase a valid key. + indices = mx.where(selected_slot < counts[..., None], indices, n_tiles) + keep = mx.zeros((heads, count, n_tiles + 1), dtype=mx.bool_) + tile_index = mx.arange(count, dtype=mx.int32)[None, :, None] + keep[head_index, tile_index, indices] = True + mask = (keep[:, :, :n_tiles, None] & key_valid[None, None]).reshape(1, heads * count, 1, padded) + query = mx.contiguous(q_video[:, start:end].reshape(1, heads * count, tile, dim)) + output = mx.fast.scaled_dot_product_attention(query, + k_shared, + v_shared, + scale=scale, + mask=mask, + force_fused=True) + output = output.reshape(heads, count, tile, dim) + valid_query = row[None, :] < sizes[n_prefix + start:n_prefix + end, None] + output = mx.where(valid_query[None, :, :, None] & (counts[..., None, None] > 0), output, 0) + mx.async_eval(output) + outputs.append(output) + if len(outputs) > 1: + mx.eval(outputs[-2]) + if not outputs: + return mx.zeros_like(q_tiled) + mx.eval(outputs[-1]) + video = mx.concatenate(outputs, axis=1) if len(outputs) > 1 else outputs[0] + video = video.transpose(1, 2, 0, 3).reshape(n_video * tile, heads, dim) + return mx.concatenate((mx.zeros((n_prefix * tile, heads, dim), dtype=q_tiled.dtype), video), axis=0) diff --git a/fastvideo/mlx_runtime/minimax_h3_vsa_target.py b/fastvideo/mlx_runtime/minimax_h3_vsa_target.py index 50eea2d89..3f30576c3 100644 --- a/fastvideo/mlx_runtime/minimax_h3_vsa_target.py +++ b/fastvideo/mlx_runtime/minimax_h3_vsa_target.py @@ -43,6 +43,9 @@ def h3_vsa_attention_target( ref_keep_rate: float | None = None, ): """query/gate_compress: [T, H, D] target rows; key/value: [S, H, D] packed. Returns [T, H, D].""" + if impl in ("sdpa", "nax"): + raise ValueError( + f"{impl.upper()} VSA does not support prefix K/V reuse; disable reuse or select another backend") import mlx.core as mx _, heads, dim = query.shape diff --git a/fastvideo/mlx_runtime/tests/test_minimax_h3_text_refiner_cache.py b/fastvideo/mlx_runtime/tests/test_minimax_h3_text_refiner_cache.py new file mode 100644 index 000000000..eb0584784 --- /dev/null +++ b/fastvideo/mlx_runtime/tests/test_minimax_h3_text_refiner_cache.py @@ -0,0 +1,120 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Request-local text refinement must preserve denoising and refresh per request.""" +from types import SimpleNamespace + +import numpy as np +import pytest + +mx = pytest.importorskip("mlx.core") + +from fastvideo.mlx_runtime import minimax_h3 as h3 +from fastvideo.mlx_runtime import minimax_h3_pipeline as pipeline_module + + +def _toy_dit(): + rng = np.random.default_rng(11) + + def array(shape): + return mx.array(rng.normal(0, 0.05, shape).astype(np.float32)) + + hidden, ffn = 16, 24 + config = dict(hidden_size=hidden, num_attention_heads=1, attention_head_dim=hidden, + ffn_dim=ffn, in_channels=4, audio_in_channels=4, patch_size=[1, 1, 1], + text_dim=8, freq_dim=hidden, time_embed_dim=hidden, rope_freq_dim=2, + rope_theta=10000., norm_eps=1e-5, qk_norm_eps=1e-5, final_norm_eps=1e-5) + weights = {} + for name, source, target in (("context_embedder", 8, hidden), ("proj_in", 4, hidden), + ("audio_proj_in", 4, hidden), ("proj_out", hidden, 4), + ("audio_proj_out", hidden, 4), + ("time_embedder.linear_1", hidden, hidden), + ("time_embedder.linear_2", hidden, hidden), + ("norm_out.linear", hidden, 2 * hidden)): + weights[name + ".weight"] = array((target, source)) + weights[name + ".bias"] = array((target,)) + weights["token_refiner.final_norm.weight"] = mx.ones((hidden,)) + weights["norm_out.norm.weight"] = mx.ones((hidden,)) + + def block(): + result = {name: mx.ones((hidden,)) for name in + ("norm1.weight", "norm2.weight", "attn.norm_q.weight", "attn.norm_k.weight")} + for name in ("attn.to_q.weight", "attn.to_k.weight", "attn.to_v.weight", "attn.to_out.0.weight"): + result[name] = array((hidden, hidden)) + result["ff.net.0.proj.weight"] = array((2 * ffn, hidden)) + result["ff.net.2.weight"] = array((hidden, ffn)) + result["adaln_proj.linear.weight"] = array((18 * hidden, hidden)) + result["adaln_proj.linear.bias"] = array((18 * hidden,)) + return result + + return h3.MLXMiniMaxH3DiT(weights, [block()], [block()], config) + + +def test_precomputed_refiner_is_exact_and_runs_once_for_eight_steps(monkeypatch): + dit = _toy_dit() + layout = h3.build_packed_layout(3, 1, 2, 2, 2, patch_size=dit.patch_size) + timesteps, inverse = h3.build_row_timesteps(layout, .7, .4) + dit.precompute_adaln(timesteps) + text = mx.array(np.arange(24, dtype=np.float32).reshape(3, 8) / 24) + video = mx.ones((len(layout.video_indices), 4)) + audio = mx.ones((len(layout.audio_indices), 4)) + calls = [] + original = dit.refine_text + + def counted(rows): + calls.append(rows) + return original(rows) + + monkeypatch.setattr(dit, "refine_text", counted) + kwargs = dict(layout=layout, step_timesteps=timesteps, row_timestep_inverse=inverse) + expected = [dit.forward_with_cache(video * (step + 1), audio, text, **kwargs) for step in range(8)] + mx.eval(expected) + assert len(calls) == 8 + calls.clear() + embeds = dit.refine_text(text) + mx.eval(embeds) + actual = [dit.forward_with_cache(video * (step + 1), audio, text, + precomputed_text_embeds=embeds, **kwargs) for step in range(8)] + mx.eval(actual) + assert len(calls) == 1 + for reference, cached in zip(expected, actual, strict=True): + for left, right in zip(reference, cached, strict=True): + np.testing.assert_array_equal(np.asarray(left), np.asarray(right)) + with pytest.raises(ValueError, match="shape"): + dit.forward_with_cache(video, audio, text, precomputed_text_embeds=embeds[:1], **kwargs) + with pytest.raises(ValueError, match="floating-point"): + dit.forward_with_cache(video, audio, text, precomputed_text_embeds=embeds.astype(mx.int32), **kwargs) + + +def test_pipeline_refines_again_for_next_request(monkeypatch, tmp_path): + dit = _toy_dit() + video_schedule = h3.MiniMaxH3SchedulerState.create(10, 8) + audio_schedule = h3.MiniMaxH3SchedulerState.create(3, 8) + union = h3.adaln_timestep_union(video_schedule, audio_schedule) + schedule = SimpleNamespace(video=video_schedule, audio=audio_schedule, adaln_timesteps=union, + num_steps=8, pdd_plan=None) + monkeypatch.setattr(pipeline_module, "resolve_h3_denoise_schedule", lambda *a: schedule) + monkeypatch.setattr(pipeline_module, "_cleanup_mlx", lambda: None) + pipeline = object.__new__(pipeline_module.MiniMaxH3MLXPipeline) + pipeline.resident = False + pipeline.model_root = tmp_path + pipeline.resolve_geometry = lambda *a, **k: dict(num_frames=124, latent_frame_count=1, + latent_height=2, latent_width=2) + calls = [] + original = dit.refine_text + + def counted(rows): + calls.append(np.asarray(rows)) + return original(rows) + + monkeypatch.setattr(dit, "refine_text", counted) + first_text = np.arange(24, dtype=np.float32).reshape(3, 8) / 24 + second_text = first_text[:, ::-1].copy() + kwargs = dict(height=32, width=32, num_frames=124, seed=0, num_steps=8, dit=dit) + first = pipeline.denoise(first_text, np.ones(3, dtype=np.int64), **kwargs) + second = pipeline.denoise(second_text, np.ones(3, dtype=np.int64), **kwargs) + repeated = pipeline.denoise(first_text, np.ones(3, dtype=np.int64), **kwargs) + assert len(calls) == 3 + np.testing.assert_array_equal(calls[0], first_text) + np.testing.assert_array_equal(calls[1], second_text) + assert not np.array_equal(first[0], second[0]) + np.testing.assert_array_equal(first[0], repeated[0]) + np.testing.assert_array_equal(first[1], repeated[1]) diff --git a/fastvideo/mlx_runtime/tests/test_minimax_h3_vsa_gather_budget.py b/fastvideo/mlx_runtime/tests/test_minimax_h3_vsa_gather_budget.py index 59fabf5f0..0c83827bd 100644 --- a/fastvideo/mlx_runtime/tests/test_minimax_h3_vsa_gather_budget.py +++ b/fastvideo/mlx_runtime/tests/test_minimax_h3_vsa_gather_budget.py @@ -50,5 +50,10 @@ def test_compiled_rotary_matches_uncompiled_fp32_math(): rotated = mx.concatenate([-second, first], axis=-1) out = (x_rot * cos[:, None, :] + rotated * sin[:, None, :]).astype(x.dtype) expected = mx.concatenate([out, x[..., rotary_dim:]], axis=-1) - np.testing.assert_array_equal(np.array(apply_h3_rotary(x, cos, sin).astype(mx.float32)), - np.array(expected.astype(mx.float32))) + # The compiled kernel may fuse the multiply-add (FMA) where the eager graph rounds each op, + # which can move a result by one bf16 ULP on some CPUs. Allow exactly that: rtol 2**-7 is one + # bf16 ULP at the bottom of each binade, so a real math error still fails. + np.testing.assert_allclose(np.array(apply_h3_rotary(x, cos, sin).astype(mx.float32)), + np.array(expected.astype(mx.float32)), + rtol=2**-7, + atol=0) diff --git a/fastvideo/tests/mlx/test_mlx_minimax_h3_prefix_kv.py b/fastvideo/tests/mlx/test_mlx_minimax_h3_prefix_kv.py index f26ed73b5..500b6ce63 100644 --- a/fastvideo/tests/mlx/test_mlx_minimax_h3_prefix_kv.py +++ b/fastvideo/tests/mlx/test_mlx_minimax_h3_prefix_kv.py @@ -80,8 +80,10 @@ def _cache_for(dit, layout, storage="native"): return MiniMaxH3PrefixKVCache(len(dit.blocks), frozen_prefix_rows(layout), dit.num_heads, dit.head_dim, storage) -def _full(dit, layout, inputs, record=None, video_t=0.7, audio_t=0.4, step_index=0): +def _full(dit, layout, inputs, record=None, video_t=0.7, audio_t=0.4, step_index=0, text_embeds=None): kwargs = {"record_prefix_kv": record} if record is not None else {} + if text_embeds is not None: + kwargs["precomputed_text_embeds"] = text_embeds video, audio = dit.forward_with_cache(inputs["x_v"], inputs["x_a"], inputs["text"], layout=layout, step_index=step_index, **_step_args(layout, video_t, audio_t), **kwargs) mx.eval(video, audio) @@ -112,6 +114,24 @@ def test_recording_step_matches_the_normal_forward_exactly(tiny_dit) -> None: assert cache.nbytes == expected +def test_precomputed_text_preserves_recorded_prefix_kv(tiny_dit) -> None: + """Cached refinement must preserve both visible outputs and later-step K/V.""" + dit, layout = tiny_dit + inputs = _inputs(layout, seed=13) + expected_cache = _cache_for(dit, layout) + expected = _full(dit, layout, inputs, record=expected_cache) + text_embeds = dit.refine_text(inputs["text"]) + mx.eval(text_embeds) + actual_cache = _cache_for(dit, layout) + actual = _full(dit, layout, inputs, record=actual_cache, text_embeds=text_embeds) + for left, right in zip(expected, actual, strict=True): + np.testing.assert_array_equal(left, right) + assert actual_cache.complete + for layer in range(len(dit.blocks)): + for left, right in zip(expected_cache.load(layer), actual_cache.load(layer), strict=True): + np.testing.assert_array_equal(np.asarray(left), np.asarray(right)) + + def test_reuse_with_unmoved_inputs_equals_the_full_forward(tiny_dit) -> None: """If nothing the prefix depends on changed, frozen K/V is exactly what the full pass computes.""" dit, layout = tiny_dit @@ -290,13 +310,20 @@ class _StubDiT: def __init__(self) -> None: self.calls: list[str] = [] + self.refiner_inputs: list[np.ndarray] = [] + self.text_embeddings: list[object] = [] def precompute_adaln(self, timesteps, *, drop_weights: bool = True): return None + def refine_text(self, text): + self.refiner_inputs.append(np.asarray(text).copy()) + return text + 1.0 + def forward_with_cache(self, x_v, x_a, text, *, layout, step_timesteps, row_timestep_inverse, step_index=0, - record_prefix_kv=None): + record_prefix_kv=None, precomputed_text_embeds=None): self.calls.append("record" if record_prefix_kv is not None else "full") + self.text_embeddings.append(precomputed_text_embeds) return x_v * 0.0 + 1.0, x_a * 0.0 + 1.0 def forward_target_rows(self, x_v, x_a, *, layout, step_timesteps, row_timestep_inverse, prefix_kv, step_index=0): @@ -336,10 +363,18 @@ def test_denoise_loop_records_then_reuses_and_refreshes(tmp_path: Path) -> None: stats = pipeline.last_prefix_kv_stats assert stats["enabled"] and stats["recorded_steps"] == [0] and stats["reused_steps"] == [1, 2, 3] assert len(pipeline.last_denoise_step_times) == 4 + assert len(dit.refiner_inputs) == len(dit.text_embeddings) == 1 + np.testing.assert_array_equal(np.asarray(dit.text_embeddings[0]), dit.refiner_inputs[0] + 1.0) dit = _StubDiT() pipeline = _denoise(tmp_path, dit, num_steps=4, reuse_prefix_kv=True, prefix_kv_refresh_every=2) assert dit.calls == ["record", "reuse", "record", "reuse"] + assert pipeline.last_prefix_kv_stats["recorded_steps"] == [0, 2] + assert len(dit.refiner_inputs) == 1 + assert len(dit.text_embeddings) == 2 + # Refreshing transformer K/V does not invalidate timestep-independent text. + for embeds in dit.text_embeddings: + np.testing.assert_array_equal(np.asarray(embeds), dit.refiner_inputs[0] + 1.0) def test_denoise_default_never_touches_prefix_kv(tmp_path: Path) -> None: @@ -347,6 +382,10 @@ def test_denoise_default_never_touches_prefix_kv(tmp_path: Path) -> None: pipeline = _denoise(tmp_path, dit, num_steps=2) assert dit.calls == ["full", "full"] assert pipeline.last_prefix_kv_stats["enabled"] is False + assert len(dit.refiner_inputs) == 1 + assert len(dit.text_embeddings) == 2 + for embeds in dit.text_embeddings: + np.testing.assert_array_equal(np.asarray(embeds), dit.refiner_inputs[0] + 1.0) def test_denoise_refuses_a_cache_over_budget(tmp_path: Path) -> None: @@ -354,6 +393,7 @@ def test_denoise_refuses_a_cache_over_budget(tmp_path: Path) -> None: with pytest.raises(PrefixKVBudgetError, match="prefix_kv_cache_max_gib"): _denoise(tmp_path, dit, reuse_prefix_kv=True, prefix_kv_cache_max_gib=1e-9) assert dit.calls == [] + assert dit.refiner_inputs == [] def test_reuse_needs_references() -> None: diff --git a/fastvideo/tests/mlx/test_mlx_minimax_h3_ref2va_pipeline.py b/fastvideo/tests/mlx/test_mlx_minimax_h3_ref2va_pipeline.py index dffc5979f..0b7d5fcdf 100644 --- a/fastvideo/tests/mlx/test_mlx_minimax_h3_ref2va_pipeline.py +++ b/fastvideo/tests/mlx/test_mlx_minimax_h3_ref2va_pipeline.py @@ -71,14 +71,22 @@ class _RecordingDiT: def __init__(self) -> None: self.calls: list[dict[str, np.ndarray]] = [] + self.refiner_inputs: list[np.ndarray] = [] + self.text_embeddings: list[object] = [] def precompute_adaln(self, timesteps, *, drop_weights: bool = True): self.adaln_timesteps = np.asarray(timesteps) return None - def forward_with_cache(self, x_v, x_a, text, *, layout, step_timesteps, row_timestep_inverse, step_index=0): - import mlx.core as mx + def refine_text(self, text): + self.refiner_inputs.append(np.asarray(text).copy()) + # Distinguishable from the raw conditioning so a missing cache argument + # cannot accidentally satisfy the pipeline handoff assertions. + return text * 2.0 + 1.0 + def forward_with_cache(self, x_v, x_a, text, *, layout, step_timesteps, row_timestep_inverse, step_index=0, + precomputed_text_embeds=None): + self.text_embeddings.append(precomputed_text_embeds) self.calls.append({ "video_rows": np.asarray(x_v), "audio_rows": np.asarray(x_a), @@ -180,9 +188,10 @@ def test_denoise_hands_the_encoder_rows_to_every_forward(tmp_path: Path) -> None video_conditions = np.full((layout.num_condition_video_rows, PATCH_DIM), 3.0, dtype=np.float32) audio_conditions = np.full((layout.num_condition_audio_rows, AUDIO_IN_CHANNELS), 4.0, dtype=np.float32) dit = _RecordingDiT() + text_rows = np.arange(NUM_TEXT_TOKENS * 16, dtype=np.float32).reshape(NUM_TEXT_TOKENS, 16) / 16 video_rows, audio_rows = pipeline.denoise( - np.zeros((NUM_TEXT_TOKENS, 16), dtype=np.float32), + text_rows, np.full(NUM_TEXT_TOKENS, 1, dtype=np.int64), height=HEIGHT, width=WIDTH, @@ -197,6 +206,11 @@ def test_denoise_hands_the_encoder_rows_to_every_forward(tmp_path: Path) -> None ) assert len(dit.calls) == 2 + assert len(dit.refiner_inputs) == 1 + np.testing.assert_array_equal(dit.refiner_inputs[0], text_rows) + assert len(dit.text_embeddings) == len(dit.calls) + for embeds in dit.text_embeddings: + np.testing.assert_array_equal(np.asarray(embeds), text_rows * 2.0 + 1.0) for call in dit.calls: assert call["video_rows"].shape == (layout.video_indices.shape[0], PATCH_DIM) assert call["audio_rows"].shape == (layout.audio_indices.shape[0], AUDIO_IN_CHANNELS) diff --git a/fastvideo/tests/mlx/test_mlx_minimax_h3_vsa.py b/fastvideo/tests/mlx/test_mlx_minimax_h3_vsa.py index e8043f1e6..ac44089dc 100644 --- a/fastvideo/tests/mlx/test_mlx_minimax_h3_vsa.py +++ b/fastvideo/tests/mlx/test_mlx_minimax_h3_vsa.py @@ -387,7 +387,7 @@ def test_untile_hidden_never_consumes_padded_query_rows() -> None: def test_simd_unsupported_tile_falls_back_to_reference() -> None: - assert PUBLIC_VSA_IMPLS == ("auto", "reference", "simd") + assert PUBLIC_VSA_IMPLS == ("auto", "reference", "simd", "sdpa", "nax") assert "metal" not in PUBLIC_VSA_IMPLS assert resolve_impl("simd", head_dim=128, tile_elems=256) == "reference" assert resolve_impl("simd", head_dim=64, tile_elems=64) == "reference" diff --git a/fastvideo/tests/mlx/test_mlx_minimax_h3_vsa_nax.py b/fastvideo/tests/mlx/test_mlx_minimax_h3_vsa_nax.py new file mode 100644 index 000000000..8bf2011bf --- /dev/null +++ b/fastvideo/tests/mlx/test_mlx_minimax_h3_vsa_nax.py @@ -0,0 +1,373 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Independent selected-token checks for the explicitly requested M5 backend.""" +from __future__ import annotations + +import hashlib +import json +from pathlib import PurePosixPath + +import numpy as np +import pytest + +mx = pytest.importorskip("mlx.core", reason="Apple MLX is required") + +from fastvideo.mlx_runtime import minimax_h3_vsa as vsa +from fastvideo.mlx_runtime import minimax_h3_vsa_nax as nax + + +@pytest.fixture +def nax_device(): + if not mx.metal.is_available(): + pytest.skip("NAX attention requires Metal on an Apple M5 GPU") + if not str(mx.device_info().get("device_name", "")).startswith("Apple M5"): + pytest.skip("NAX attention is validated only on Apple M5 GPUs") + + +def _selected_rows(indices, counts, geometry, head, query): + tile = geometry.tile_elems + query_tile = geometry.num_prefix_tiles + query + query_rows = np.arange(query_tile * tile, query_tile * tile + geometry.variable_block_sizes[query_tile]) + selected = indices[head, query, :counts[head, query]] + key_rows = (np.concatenate([np.arange(index * tile, index * tile + geometry.variable_block_sizes[index]) + for index in selected]) if len(selected) else np.empty(0, dtype=np.int64)) + return query_rows, key_rows + + +def _numpy_reference(q, k, value, indices, counts, geometry, scale): + """Select individual valid tokens independently of all MLX block masks.""" + result = np.zeros_like(q, dtype=np.float64) + for head in range(q.shape[1]): + for query in range(geometry.num_video_tiles): + query_rows, key_rows = _selected_rows(indices, counts, geometry, head, query) + if not len(key_rows): + continue + scores = q[query_rows, head].astype(np.float64) @ k[key_rows, head].astype(np.float64).T * scale + probabilities = np.exp(scores - scores.max(axis=-1, keepdims=True)) + probabilities /= probabilities.sum(axis=-1, keepdims=True) + result[query_rows, head] = probabilities @ value[key_rows, head].astype(np.float64) + return result + + +def _native_reference(q, k, value, indices, counts, geometry, scale): + """Gather the same valid tokens; no custom kernel, padding, masks or GQA.""" + result = np.zeros(q.shape, dtype=np.float32) + for head in range(q.shape[1]): + for query in range(geometry.num_video_tiles): + query_rows, key_rows = _selected_rows(indices, counts, geometry, head, query) + if not len(key_rows): + continue + output = mx.fast.scaled_dot_product_attention( + q[mx.array(query_rows), head][None, None], + k[mx.array(key_rows), head][None, None], + value[mx.array(key_rows), head][None, None], scale=scale, force_fused=True) + mx.eval(output) + result[query_rows, head] = np.asarray(output[0, 0].astype(mx.float32)) + return result + + +def _metrics(actual, expected): + difference = actual.astype(np.float64) - expected + return {"max_abs": float(np.abs(difference).max()), + "rmse": float(np.sqrt(np.mean(difference**2))), + "relative_l2": float(np.linalg.norm(difference) / max(np.linalg.norm(expected), 1e-30))} + + +def _case(tile_size, dtype_name, *, segments=None): + # Dense/reference/target regions are interleaved before tiling. Ragged tiles + # cross 16/32/64-row boundaries, including inactive SIMD groups in a launch. + geometry = vsa.build_h3_segment_geometry(segments or (7, (1, 3, 5), 13, (3, 7, 11)), tile_size) + rng = np.random.default_rng(9241) + shape = (geometry.padded_length, 3, 128) + arrays = [rng.normal(size=shape).astype(np.float32) * 0.6 for _ in range(3)] + padding = np.concatenate([np.arange(index * tile_size + size, (index + 1) * tile_size) + for index, size in enumerate(geometry.variable_block_sizes)]) + for index, array in enumerate(arrays): + array[padding] = 30.0 + index + width = min(5, geometry.num_tiles) + # Intentionally out of bounds: inactive route slots must never be read. + indices = np.full((3, geometry.num_video_tiles, width), np.iinfo(np.int32).max, dtype=np.int32) + counts = np.empty((3, geometry.num_video_tiles), dtype=np.int32) + for head in range(3): + for query in range(geometry.num_video_tiles): + count = 1 + (head + query) % width + counts[head, query] = count + indices[head, query, :count] = rng.permutation(geometry.num_tiles)[:count] + counts[-1, -1] = 0 + indices[-1, -1] = np.iinfo(np.int32).max + tensors = [mx.array(array).astype(getattr(mx, dtype_name)) for array in arrays] + mx.eval(*tensors) + return geometry, tensors, indices, counts, padding + + +# MLX's installed steel/attn/nax.h sets MPP relaxed_precision=true even for +# float32. The native gathered control reproduces its ~7e-4 NumPy discrepancy +# on M5. Keep the 1e-3 independent FP32 bound and record BOTH numerical controls; +# changing the selected tokens must not be excused as a precision difference. +# The probability accumulation schedule differs from the native dense kernel. +@pytest.mark.parametrize("dtype_name,tolerance", [("float16", 0.003), ("bfloat16", 0.02), ("float32", 0.001)]) +@pytest.mark.parametrize("tile_size,query_rows", [(128, 16), (128, 32), (128, 64), (128, 128), + (64, 16), (64, 32), (64, 64)]) +def test_nax_matches_independent_valid_token_oracle(nax_device, record_property, dtype_name, + tolerance, tile_size, query_rows): + geometry, tensors, indices, counts, padding = _case(tile_size, dtype_name) + scale = 128**-0.5 + rounded = [np.asarray(tensor.astype(mx.float32)) for tensor in tensors] + expected = _numpy_reference(*rounded, indices, counts, geometry, scale) + actual = nax.nax_block_sparse(*tensors, mx.array(indices), mx.array(counts), geometry, + scale, query_rows=query_rows) + mx.eval(actual) + got = np.asarray(actual.astype(mx.float32)) + record_property("candidate_vs_numpy", json.dumps(_metrics(got, expected))) + assert actual.shape == tensors[0].shape + assert actual.dtype == tensors[0].dtype + assert np.isfinite(got).all() + np.testing.assert_allclose(got, expected, atol=tolerance, rtol=tolerance) + np.testing.assert_array_equal(got[padding], 0) + np.testing.assert_array_equal(got[:geometry.num_prefix_tiles * tile_size], 0) + zero_rows, _ = _selected_rows(indices, counts, geometry, 2, geometry.num_video_tiles - 1) + np.testing.assert_array_equal(got[zero_rows, 2], 0) + if dtype_name == "float32": + native = _native_reference(*tensors, indices, counts, geometry, scale) + record_property("native_vs_numpy", json.dumps(_metrics(native, expected))) + record_property("candidate_vs_native", json.dumps(_metrics(got, native))) + np.testing.assert_allclose(native, expected, atol=tolerance, rtol=tolerance) + # Both paths have independently bounded NumPy error. This pairwise + # budget is their sum, not a claim of bitwise native-kernel equivalence. + np.testing.assert_allclose(got, native, atol=2 * tolerance, rtol=2 * tolerance) + + +@pytest.mark.parametrize("query_rows", [16, 32, 64, 128]) +def test_nax_uniform_logits_expose_head_and_tile_selection(nax_device, query_rows): + geometry, tensors, indices, counts, padding = _case(128, "float32") + shape = tensors[0].shape + q, k = mx.zeros(shape), mx.zeros(shape) + value = np.zeros(shape, dtype=np.float32) + for tile, size in enumerate(geometry.variable_block_sizes): + for head in range(shape[1]): + # Each tile and head carries a distinct basis vector. Uniform + # logits make the answer a valid-row-weighted tile histogram. + value[tile * 128:tile * 128 + size, head, tile] = head + 1 + value[padding] = 32.0 + expected = _numpy_reference(np.zeros(shape), np.zeros(shape), value, indices, counts, geometry, 0.071) + actual = nax.nax_block_sparse(q, k, mx.array(value), mx.array(indices), mx.array(counts), + geometry, 0.071, query_rows=query_rows) + mx.eval(actual) + np.testing.assert_allclose(np.asarray(actual), expected, atol=2e-6, rtol=2e-6) + + +def test_nax_without_dense_prefix_and_with_nondefault_scale(nax_device): + geometry, tensors, indices, counts, padding = _case(64, "bfloat16", segments=((3, 7, 11),)) + assert geometry.num_prefix_tiles == 0 + expected = _numpy_reference(*[np.asarray(t.astype(mx.float32)) for t in tensors], + indices, counts, geometry, 0.17) + actual = nax.nax_block_sparse(*tensors, mx.array(indices), mx.array(counts), geometry, 0.17, + query_rows=32) + mx.eval(actual) + np.testing.assert_allclose(np.asarray(actual.astype(mx.float32)), expected, atol=0.02, rtol=0.02) + np.testing.assert_array_equal(np.asarray(actual.astype(mx.float32))[padding], 0) + + +def _attention_inputs(): + geometry = vsa.build_h3_segment_geometry((7, (1, 3, 5), 13, (3, 7, 11)), 128) + rng = np.random.default_rng(874) + shape = (geometry.total_seq_length, 2, 128) + tensors = [mx.array(rng.normal(size=shape).astype(np.float32) * 0.1) for _ in range(4)] + return geometry, tensors + + +def test_explicit_nax_preserves_router_dense_prefix_and_compression(nax_device, monkeypatch): + geometry, (q, k, value, gate) = _attention_inputs() + expected = vsa.h3_vsa_attention(q, k, value, geometry, sparsity=0.5, ref_keep_rate=0.5, + gate_compress=gate, impl="reference") + calls = [] + original = nax.nax_block_sparse + + def record(*args, **kwargs): + calls.append((args[3].shape, args[4].shape)) + return original(*args, **kwargs) + + monkeypatch.setattr(nax, "nax_block_sparse", record) + stats = vsa.MiniMaxH3VSAStats() + actual = vsa.h3_vsa_attention(q, k, value, geometry, sparsity=0.5, ref_keep_rate=0.5, + gate_compress=gate, impl="nax", stats=stats) + mx.eval(actual, expected) + assert calls == [((2, geometry.num_video_tiles, calls[0][0][-1]), (2, geometry.num_video_tiles))] + assert stats.impl == "nax" + assert stats.dense_fallback_reason is None + np.testing.assert_allclose(np.asarray(actual), np.asarray(expected), atol=2e-5, rtol=2e-4) + + +def test_explicit_nax_propagates_kernel_failure(nax_device, monkeypatch): + geometry, (q, k, value, _) = _attention_inputs() + + def fail(*args, **kwargs): + raise RuntimeError("deliberate NAX kernel failure") + + monkeypatch.setattr(nax, "nax_block_sparse", fail) + with pytest.raises(RuntimeError, match="deliberate NAX kernel failure"): + actual = vsa.h3_vsa_attention(q, k, value, geometry, sparsity=0.5, ref_keep_rate=0.5, impl="nax") + mx.eval(actual) + + +def test_nax_requires_explicit_selection(monkeypatch): + from fastvideo.mlx_runtime import minimax_h3_vsa_simd as simd + + monkeypatch.setattr(simd, "simd_kernel_available", lambda: True) + assert "nax" in vsa.PUBLIC_VSA_IMPLS + assert vsa.resolve_impl("nax", 128, 128) == "nax" + assert vsa.resolve_impl("nax", 128, 64) == "nax" + assert vsa.resolve_impl("auto", 128, 128) == "simd" + assert vsa.resolve_impl("reference", 128, 128) == "reference" + monkeypatch.setattr(simd, "simd_kernel_available", lambda: False) + monkeypatch.setattr(simd, "simd_kernel_error", lambda: "test unavailable") + assert vsa.resolve_impl("auto", 128, 128) == "reference" + + +@pytest.mark.parametrize("head_dim,tile_size", [(64, 128), (128, 256)]) +def test_nax_rejects_unsupported_public_shapes(head_dim, tile_size): + with pytest.raises(ValueError): + vsa.resolve_impl("nax", head_dim, tile_size) + + +@pytest.mark.parametrize("device_name", ["Apple M4 Max", "Apple M3", "Apple M50", "AMD Radeon", ""]) +def test_nax_rejects_unsupported_device_before_kernel_construction(monkeypatch, device_name): + monkeypatch.setattr(mx.metal, "is_available", lambda: True) + monkeypatch.setattr(mx, "device_info", lambda: {"device_name": device_name}, raising=False) + + def unexpected(*args, **kwargs): + pytest.fail("An unsupported device must not construct a Metal kernel") + + monkeypatch.setattr(mx.fast, "metal_kernel", unexpected, raising=False) + nax._nax_kernel.cache_clear() + try: + with pytest.raises(RuntimeError, match="M5"): + nax._nax_kernel() + finally: + nax._nax_kernel.cache_clear() + + +def test_nax_rejects_missing_metal(monkeypatch): + monkeypatch.setattr(mx.metal, "is_available", lambda: False) + with pytest.raises(RuntimeError, match="Metal|M5"): + nax._check_nax_device() + + +def test_nax_construction_failure_keeps_original_cause(monkeypatch): + original_error = RuntimeError("deliberate compiler error") + monkeypatch.setattr(nax, "_check_nax_device", lambda: None) + monkeypatch.setattr(nax, "_header", lambda: ("unused", {})) + + def fail(*args, **kwargs): + raise original_error + + monkeypatch.setattr(mx.fast, "metal_kernel", fail, raising=False) + nax._nax_kernel.cache_clear() + try: + with pytest.raises(RuntimeError, match="kernel construction") as captured: + nax._nax_kernel() + assert captured.value.__cause__ is original_error + finally: + nax._nax_kernel.cache_clear() + + +def test_unvalidated_precision_override_is_not_a_public_option(): + with pytest.raises(TypeError, match="relaxed_precision"): + nax.nax_block_sparse(None, None, None, None, None, None, 1.0, + relaxed_precision=False) + + +def _assert_nax_provenance_schema(provenance): + assert provenance["relaxed_precision"] is True + assert provenance["math_mode"] == "safe" + assert provenance["key_rows"] == 32 + for name in ("source_sha256", "expanded_header_sha256"): + assert len(provenance[name]) == 64 + int(provenance[name], 16) + dependencies = provenance["installed_header_sha256"] + assert dependencies + for path, digest in dependencies.items(): + assert path.startswith("mlx/") + assert not PurePosixPath(path).is_absolute() + assert ".." not in PurePosixPath(path).parts + assert len(digest) == 64 + int(digest, 16) + + +def test_nax_provenance_schema_without_installed_metal_headers(monkeypatch): + # The Linux MLX CPU wheel need not ship Apple's Metal helper headers. + # Exercise receipt formatting and hashing without consulting that package. + header = "// synthetic expanded NAX header\n" + dependencies = {"mlx/backend/metal/kernels/steel/attn/nax.h": hashlib.sha256(b"synthetic header").hexdigest()} + monkeypatch.setattr(nax, "_header", lambda: (header, dependencies)) + provenance = nax.get_nax_provenance() + _assert_nax_provenance_schema(provenance) + assert provenance["expanded_header_sha256"] == hashlib.sha256(header.encode()).hexdigest() + assert provenance["source_sha256"] == hashlib.sha256(nax._SOURCE.encode()).hexdigest() + assert provenance["installed_header_sha256"] == dependencies + + +def test_nax_installed_header_provenance_on_supported_device(nax_device): + # A supported M5 installation must supply the real validated headers; + # missing or incompatible headers are a failure here, not a skip. + provenance = nax.get_nax_provenance() + _assert_nax_provenance_schema(provenance) + include_root = nax._mlx_include_root() + for relative, digest in provenance["installed_header_sha256"].items(): + assert hashlib.sha256((include_root / relative).read_bytes()).hexdigest() == digest + + +@pytest.fixture +def synthetic_nax_headers(monkeypatch, tmp_path): + # Parse a tiny include tree without Metal or the installed private headers. + # The expected API markers deliberately match the loader's known contract. + dependency = "mlx/testing/nax_dependency.h" + descriptor = "transpose_b,\n true,\n mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate\n" + header = (f'#include "{dependency}"\n' + "struct BaseNAXFrag; struct NAXTile;\n" + "get_coord(); row_reduce(); row_bin_op(); load_rows(); store_rows();\n" + + descriptor * 2) + files = {nax._NAX_HEADER: header, dependency: "// validated dependency\n"} + hashes = {} + for relative, text in files.items(): + path = tmp_path / relative + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(text) + hashes[relative] = hashlib.sha256(path.read_bytes()).hexdigest() + monkeypatch.setattr(nax, "_mlx_include_root", lambda: tmp_path) + monkeypatch.setattr(nax, "_VALIDATED_HEADER_SHA256", hashes) + nax._header.cache_clear() + try: + yield tmp_path, dependency, hashes + finally: + nax._header.cache_clear() + + +def test_nax_rejects_modified_validated_header(synthetic_nax_headers): + include_root, dependency, hashes = synthetic_nax_headers + _, actual_hashes = nax._header() + assert actual_hashes == hashes + (include_root / dependency).write_text("// changed dependency with unchanged API\n") + nax._header.cache_clear() + with pytest.raises(RuntimeError, match="Unvalidated MLX private header") as captured: + nax._header() + assert dependency in str(captured.value) + + +def test_nax_rejects_unregistered_header_dependency(synthetic_nax_headers): + include_root, _, hashes = synthetic_nax_headers + unknown = "mlx/testing/unregistered_dependency.h" + (include_root / unknown).write_text("// This dependency is outside the validated set.\n") + header_path = include_root / nax._NAX_HEADER + header_path.write_text(header_path.read_text() + f'#include "{unknown}"\n') + hashes[nax._NAX_HEADER] = hashlib.sha256(header_path.read_bytes()).hexdigest() + with pytest.raises(RuntimeError, match="Unvalidated MLX private header") as captured: + nax._header() + assert unknown in str(captured.value) + + +def test_pipeline_rejects_nax_prefix_reuse_before_loading_components(): + from fastvideo.mlx_runtime.minimax_h3_pipeline import MiniMaxH3MLXPipeline + + pipeline = object.__new__(MiniMaxH3MLXPipeline) + with pytest.raises(ValueError, match="prefix|Prefix"): + pipeline.generate("unused", output_path="unused.mp4", vsa_impl="nax", reuse_prefix_kv=True) diff --git a/fastvideo/tests/mlx/test_mlx_minimax_h3_vsa_sdpa.py b/fastvideo/tests/mlx/test_mlx_minimax_h3_vsa_sdpa.py new file mode 100644 index 000000000..1e5a12f3d --- /dev/null +++ b/fastvideo/tests/mlx/test_mlx_minimax_h3_vsa_sdpa.py @@ -0,0 +1,236 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Masked SDPA keeps H3's selected tiles and ragged-row semantics.""" +from __future__ import annotations + +import numpy as np +import pytest + +mx = pytest.importorskip("mlx.core", reason="Apple MLX is required") + +from fastvideo.mlx_runtime import minimax_h3_vsa as vsa +from fastvideo.mlx_runtime import minimax_h3_vsa_sdpa as sdpa + +pytestmark = pytest.mark.skipif(not mx.metal.is_available(), reason="Fused MLX SDPA requires Metal") + + +def _selected_token_reference(q, k, value, indices, counts, geometry): + """Independent token selection, followed by NumPy float64 softmax.""" + tile = geometry.tile_elems + sizes = geometry.variable_block_sizes + output = np.zeros_like(q, dtype=np.float64) + for head in range(q.shape[1]): + for query in range(geometry.num_video_tiles): + query_tile = geometry.num_prefix_tiles + query + query_rows = np.arange(query_tile * tile, query_tile * tile + sizes[query_tile]) + selected = indices[head, query, :counts[head, query]] + if not len(selected): + continue + key_rows = np.concatenate([np.arange(index * tile, index * tile + sizes[index]) for index in selected]) + logits = q[query_rows, head].astype(np.float64) @ k[key_rows, head].astype(np.float64).T + logits *= q.shape[-1]**-0.5 + probabilities = np.exp(logits - logits.max(axis=-1, keepdims=True)) + probabilities /= probabilities.sum(axis=-1, keepdims=True) + output[query_rows, head] = probabilities @ value[key_rows, head].astype(np.float64) + return output + + +def _gathered_native_reference(q, k, value, indices, counts, geometry): + """Native control with independently gathered valid keys and padded queries. + + Match the backend's query length: shorter ragged queries can select a + different native arithmetic path. Query rows attend independently, so + padding queries does not change the expected values of the valid rows. + This control uses no GQA, key padding or attention mask. + """ + output = np.zeros(q.shape, dtype=np.float32) + tile = geometry.tile_elems + sizes = geometry.variable_block_sizes + for head in range(q.shape[1]): + for query in range(geometry.num_video_tiles): + query_tile = geometry.num_prefix_tiles + query + selected = indices[head, query, :counts[head, query]] + if not len(selected): + continue + query_start = query_tile * tile + query_rows = np.arange(query_start, query_start + sizes[query_tile]) + key_rows = np.concatenate([np.arange(index * tile, index * tile + sizes[index]) for index in selected]) + result = mx.fast.scaled_dot_product_attention( + q[query_start:query_start + tile, head][None, None], + k[mx.array(key_rows), head][None, None], + value[mx.array(key_rows), head][None, None], + scale=q.shape[-1]**-0.5, force_fused=True) + mx.eval(result) + output[query_rows, head] = np.asarray(result[0, 0, :sizes[query_tile]].astype(mx.float32)) + return output + + +# MLX steel/attn/nax.h uses MPP matmul2d_descriptor(relaxed_precision=true), +# including float32 inputs. The MPP SDK explicitly permits an accuracy/performance +# tradeoff. On M5, the independent native gathered control and this backend both +# differ from NumPy by up to 7.45e-4 across these tile sizes. Matching their query +# lengths reduces native arithmetic differences without sharing key selection. +# Bound that native error separately from the stricter mask/layout control below. +@pytest.mark.parametrize("dtype_name,tolerance", [("float32", 1e-3), ("bfloat16", 0.02)]) +@pytest.mark.parametrize("chunk_size", [1, 3, 8]) +@pytest.mark.parametrize("tile_size", [64, 128]) +def test_sparse_sdpa_matches_selected_valid_tokens(dtype_name, tolerance, chunk_size, tile_size): + geometry = vsa.build_h3_segment_geometry((7, (1, 3, 5), 13, (3, 7, 11)), tile_size) + rng = np.random.default_rng(9241) + heads, dim = 3, 128 + shape = (geometry.padded_length, heads, dim) + originals = [rng.normal(size=shape).astype(np.float32) * 0.6 for _ in range(3)] + tile = geometry.tile_elems + sizes = geometry.variable_block_sizes + padding = np.concatenate([np.arange(index * tile + count, (index + 1) * tile) + for index, count in enumerate(sizes)]) + for index, array in enumerate(originals): + array[padding] = 30.0 + index # Missing padding masks cause large errors. + width = min(5, geometry.num_tiles) + indices = np.zeros((heads, geometry.num_video_tiles, width), dtype=np.int32) + counts = np.empty((heads, geometry.num_video_tiles), dtype=np.int32) + for head in range(heads): + for query in range(geometry.num_video_tiles): + count = 1 + (head + query) % width + counts[head, query] = count + indices[head, query, :count] = rng.permutation(geometry.num_tiles)[:count] + unselected = np.setdiff1d(np.arange(geometry.num_tiles), indices[head, query, :count]) + indices[head, query, count:] = unselected[-1] if len(unselected) else 0 + counts[-1, -1] = 0 + q, k, value = [mx.array(item).astype(getattr(mx, dtype_name)) for item in originals] + mx.eval(q, k, value) + rounded = [np.asarray(item.astype(mx.float32)) for item in (q, k, value)] + expected = _selected_token_reference(*rounded, indices, counts, geometry) + actual = sdpa.sdpa_block_sparse(q, k, value, mx.array(indices), mx.array(counts), geometry, + dim**-0.5, chunk_size=chunk_size) + mx.eval(actual) + got = np.asarray(actual.astype(mx.float32)) + assert np.isfinite(got).all() + np.testing.assert_allclose(got, expected, atol=tolerance, rtol=tolerance) + if dtype_name == "float32": + native = _gathered_native_reference(q, k, value, indices, counts, geometry) + np.testing.assert_allclose(native, expected, atol=tolerance, rtol=tolerance) + # This independent control removes all candidate mask construction, + # reshaping and GQA, so a wrong selected tile or head cannot hide behind + # the NAX kernel's looser float32 arithmetic. + np.testing.assert_allclose(got, native, atol=1e-4, rtol=2e-4) + np.testing.assert_array_equal(got[padding], 0) + np.testing.assert_array_equal(got[:geometry.num_prefix_tiles * tile], 0) + + +def test_mask_memory_cap_splits_dispatch_without_changing_selected_tokens(monkeypatch): + geometry = vsa.build_h3_segment_geometry((7, (1, 3, 5), 13, (3, 7, 11)), 128) + rng = np.random.default_rng(116) + heads, dim = 3, 128 + arrays = [rng.normal(size=(geometry.padded_length, heads, dim)).astype(np.float32) * 0.6 + for _ in range(3)] + padding = np.concatenate([np.arange(tile * 128 + count, (tile + 1) * 128) + for tile, count in enumerate(geometry.variable_block_sizes)]) + for index, array in enumerate(arrays): + array[padding] = 30.0 + index + indices = np.empty((heads, geometry.num_video_tiles, 3), dtype=np.int32) + for head in range(heads): + for query in range(geometry.num_video_tiles): + indices[head, query] = (head + query + np.array([0, 2, 4])) % geometry.num_tiles + counts = np.full((heads, geometry.num_video_tiles), 3, dtype=np.int32) + counts[-1, -1] = 0 + expected = _selected_token_reference(*arrays, indices, counts, geometry) + q, k, value = [mx.array(array) for array in arrays] + + # A fixed tiny budget exercises the automatic cap without allocating a + # production-sized mask. Observe actual native calls, not the chunk formula. + cap_bytes = 8192 + monkeypatch.setattr(sdpa, "_MASK_CHUNK_MAX_BYTES", cap_bytes) + dispatches = [] + native_sdpa = mx.fast.scaled_dot_product_attention + + def record_dispatch(query, key, value, **kwargs): + dispatches.append({"mask_bytes": kwargs["mask"].nbytes, + "query_heads": query.shape[1]}) + return native_sdpa(query, key, value, **kwargs) + + monkeypatch.setattr(mx.fast, "scaled_dot_product_attention", record_dispatch) + actual = sdpa.sdpa_block_sparse(q, k, value, mx.array(indices), mx.array(counts), geometry, + dim**-0.5, chunk_size=geometry.num_video_tiles) + mx.eval(actual) + assert len(dispatches) > 1 + assert all(0 < dispatch["mask_bytes"] <= cap_bytes for dispatch in dispatches) + assert sum(dispatch["query_heads"] for dispatch in dispatches) == heads * geometry.num_video_tiles + got = np.asarray(actual) + assert np.isfinite(got).all() + np.testing.assert_allclose(got, expected, atol=1e-3, rtol=1e-3) + np.testing.assert_array_equal(got[padding], 0) + np.testing.assert_array_equal(got[:geometry.num_prefix_tiles * geometry.tile_elems], 0) + + +def _attention_inputs(): + geometry = vsa.build_h3_segment_geometry((7, (1, 3, 5), 13, (3, 7, 11)), 128) + rng = np.random.default_rng(874) + shape = (geometry.total_seq_length, 2, 128) + tensors = [mx.array(rng.normal(size=shape).astype(np.float32) * 0.1) for _ in range(4)] + return geometry, tensors + + +def test_explicit_sdpa_routes_without_changing_attention_selection(monkeypatch): + geometry, (q, k, value, gate) = _attention_inputs() + expected = vsa.h3_vsa_attention(q, k, value, geometry, sparsity=0.5, ref_keep_rate=0.5, + gate_compress=gate, impl="reference") + calls = [] + original = sdpa.sdpa_block_sparse + + def record(*args, **kwargs): + calls.append((args[3].shape, args[4].shape)) + return original(*args, **kwargs) + + monkeypatch.setattr(sdpa, "sdpa_block_sparse", record) + stats = vsa.MiniMaxH3VSAStats() + actual = vsa.h3_vsa_attention(q, k, value, geometry, sparsity=0.5, ref_keep_rate=0.5, + gate_compress=gate, impl="sdpa", stats=stats) + mx.eval(actual, expected) + assert len(calls) == 1 + assert calls[0][0][:2] == (2, geometry.num_video_tiles) + assert calls[0][1] == (2, geometry.num_video_tiles) + assert stats.impl == "sdpa" + assert stats.dense_fallback_reason is None + np.testing.assert_allclose(np.asarray(actual), np.asarray(expected), atol=2e-5, rtol=2e-4) + + +def test_explicit_sdpa_propagates_kernel_failure(monkeypatch): + geometry, (q, k, value, _) = _attention_inputs() + + def fail(*args, **kwargs): + raise RuntimeError("test SDPA failure") + + monkeypatch.setattr(sdpa, "sdpa_block_sparse", fail) + with pytest.raises(RuntimeError, match="test SDPA failure"): + vsa.h3_vsa_attention(q, k, value, geometry, sparsity=0.5, ref_keep_rate=0.5, impl="sdpa") + + +def test_sdpa_is_explicit_and_auto_selection_stays_unchanged(monkeypatch): + from fastvideo.mlx_runtime import minimax_h3_vsa_simd as simd + + monkeypatch.setattr(simd, "simd_kernel_available", lambda: True) + assert vsa.resolve_impl("sdpa", 128, 128) == "sdpa" + assert vsa.resolve_impl("auto", 128, 128) == "simd" + + +@pytest.mark.parametrize("head_dim,tile_size", [(64, 128), (128, 256)]) +def test_explicit_sdpa_rejects_unsupported_shapes(head_dim, tile_size): + with pytest.raises(ValueError): + vsa.resolve_impl("sdpa", head_dim, tile_size) + + +def test_pipeline_rejects_sdpa_prefix_reuse_before_loading_components(): + from fastvideo.mlx_runtime.minimax_h3_pipeline import MiniMaxH3MLXPipeline + + pipeline = object.__new__(MiniMaxH3MLXPipeline) + with pytest.raises(ValueError, match="prefix|Prefix"): + pipeline.generate("unused", output_path="unused.mp4", vsa_impl="sdpa", reuse_prefix_kv=True) + + +def test_sparse_sdpa_rejects_invalid_chunk_size(): + geometry = vsa.build_h3_segment_geometry((3, (1, 3, 5)), 128) + array = mx.zeros((geometry.padded_length, 1, 128)) + indices = mx.zeros((1, geometry.num_video_tiles, 1), dtype=mx.int32) + counts = mx.ones((1, geometry.num_video_tiles), dtype=mx.int32) + with pytest.raises(ValueError, match="chunk"): + sdpa.sdpa_block_sparse(array, array, array, indices, counts, geometry, 128**-0.5, chunk_size=0) diff --git a/pyproject.toml b/pyproject.toml index 10c3edae9..bb7a482ef 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -305,7 +305,8 @@ skip = "./data,./wandb,apps/fastvideo_studio/package-lock.json,apps/performance_ # lowercase form silences all case variants. # "padd" is the literal value of Wan2.2-S2V's ``framepack_drop_mode`` in the # published checkpoint config; it must be spelled exactly to match. -ignore-words-list = "tread,passt,dout,padd" +# "nax" is the name of MLX's M5 TensorOps helper and the corresponding H3 backend. +ignore-words-list = "tread,passt,dout,padd,nax" [tool.ruff] # Allow lines to be as long as 120.