[perf] Cache MLX H3 text refinement and speed up attention (#2002)
Co-authored-by: Aryan Kumar <aryan5v@users.noreply.github.com>
This commit is contained in:
posiposi443-ops
co-authored by
Aryan Kumar
Aryan Kumar
parent
1a1e7216d2
commit
8db2d71ef2
@@ -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 \
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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 <typename V>
|
||||
METAL_FUNC static constexpr V apply(V x, V y) { return metal::max(x, y); }
|
||||
};
|
||||
struct SumOp {
|
||||
template <typename V>
|
||||
METAL_FUNC static constexpr V apply(V x, V y) { return x + y; }
|
||||
};
|
||||
struct MulOp {
|
||||
template <typename V>
|
||||
METAL_FUNC static constexpr V apply(V x, V y) { return x * y; }
|
||||
};
|
||||
struct ExpSubOp {
|
||||
template <typename V>
|
||||
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<float, 1, TD>;
|
||||
using stile_t = NAXTile<float, 1, TK>;
|
||||
otile_t Otile;
|
||||
Otile.clear();
|
||||
constexpr short kRowsPT = otile_t::kRowsPerThread;
|
||||
metal::vec<float, kRowsPT> max_score;
|
||||
metal::vec<float, kRowsPT> 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<T, 1, 1> Qtile;
|
||||
NAXTile<T, 2, 1> 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<float, kRowsPT> new_max = max_score;
|
||||
metal::vec<float, kRowsPT> factor;
|
||||
Stile.template row_reduce<MaxOp>(new_max);
|
||||
Stile.template row_bin_op<ExpSubOp>(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<SumOp>(sum_score);
|
||||
Otile.template row_bin_op<MulOp>(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<T, 1, 2> 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<float, kRowsPT> 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<MulOp>(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
|
||||
@@ -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)
|
||||
@@ -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
|
||||
|
||||
@@ -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])
|
||||
@@ -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)
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
+2
-1
@@ -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.
|
||||
|
||||
Reference in New Issue
Block a user