[bugfix]: scope affine dq-GEMM to H3

Author: aryan5v

.github/workflows/ci-macos-mlx.yml

@@ -81,6 +81,7 @@ jobs:
             fastvideo/tests/mlx/test_mlx_compile_parity.py \
             fastvideo/tests/mlx/test_mlx_checkpoint.py \
             fastvideo/tests/mlx/test_mlx_checkpoint_compat.py \
+            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_regressions.py \
@@ -143,6 +144,7 @@ jobs:
             fastvideo/tests/mlx/test_mlx_compile_parity.py \
             fastvideo/tests/mlx/test_mlx_checkpoint.py \
             fastvideo/tests/mlx/test_mlx_checkpoint_compat.py \
+            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_regressions.py \

fastvideo/mlx_runtime/fastwan.py

@@ -306,8 +306,10 @@ def quantize_matrix(weight, spec: MLXQuantizationSpec | None):
 # Affine quantized_matmul is slower than dequantize + steel GEMM at H3's packed
 # token width. Measured on Apple M4 Max / MLX 0.32.2, INT6 group 64, BF16 acts,
 # Q 5376→7168: M=256 qmm is faster; M=512 dequant+GEMM is +6.4%; M≥1024 ~10%.
-# 832×480×124 packed M is ~14862–14994. Do not cache dequantized weights.
-# Override: FASTVIDEO_MLX_DQ_GEMM=0 off, =1 measured floor, =<int> explicit floor.
+# 832×480×124 packed M is ~14862–14994. H3 opts into this path explicitly;
+# shared FastWan and Wan 2.2 linears stay on quantized_matmul. Do not cache
+# dequantized weights. Override: FASTVIDEO_MLX_DQ_GEMM=0 off, =1 measured
+# floor, =<int> explicit floor.
 _AFFINE_DQ_GEMM_BITS = frozenset({2, 3, 4, 5, 6, 8})
 MLX_AFFINE_DQ_GEMM_DEFAULT_MIN_M = 768
 _dq_gemm_engaged = 0
@@ -345,12 +347,12 @@ def _matmul_leading_rows(x) -> int:
     return int(x.size) // last
 
 
-def _quantized_linear(x, weight: QuantizedMatrix):
+def _quantized_linear(x, weight: QuantizedMatrix, *, use_affine_dq_gemm: bool = False):
     import mlx.core as mx
 
     global _dq_gemm_engaged, _dq_gemm_logged
     spec = weight.spec
-    min_m = affine_dq_gemm_min_m()
+    min_m = affine_dq_gemm_min_m() if use_affine_dq_gemm else None
     rows = _matmul_leading_rows(x)
     if (min_m is not None and spec.mode == "affine" and spec.bits in _AFFINE_DQ_GEMM_BITS
             and spec.group_size is not None and rows >= min_m):
@@ -381,8 +383,9 @@ def _quantized_linear(x, weight: QuantizedMatrix):
     ).astype(x.dtype)
 
 
-def linear(x, weight, bias=None):
-    y = _quantized_linear(x, weight) if isinstance(weight, QuantizedMatrix) else x @ weight.T
+def linear(x, weight, bias=None, *, use_affine_dq_gemm: bool = False):
+    y = (_quantized_linear(x, weight, use_affine_dq_gemm=use_affine_dq_gemm)
+         if isinstance(weight, QuantizedMatrix) else x @ weight.T)
     if bias is not None:
         y = y + bias
     return y

fastvideo/mlx_runtime/minimax_h3.py

@@ -61,7 +61,7 @@
     MLXQuantizationSpec,
     QuantizedMatrix,
     ensure_quantization_supported,
-    linear,
+    linear as _shared_linear,
     quantize_matrix,
     silu,
     timestep_embedding,
@@ -81,6 +81,12 @@
 
 logger = init_logger(__name__)
 
+
+def linear(x, weight, bias=None):
+    """Run an H3 linear with its measured wide-row affine dispatch enabled."""
+    return _shared_linear(x, weight, bias, use_affine_dq_gemm=True)
+
+
 # ---------------------------------------------------------------------------
 # Constants (mirrors fastvideo/pipelines/basic/minimax_h3/packing.py)
 # ---------------------------------------------------------------------------

fastvideo/tests/mlx/test_mlx_affine_dq_gemm.py

@@ -18,6 +18,7 @@
     quantize_matrix,
     reset_dq_gemm_telemetry,
 )
+from fastvideo.mlx_runtime.minimax_h3 import linear as h3_linear  # noqa: E402
 
 AFFINE_BITS = (2, 3, 4, 5, 6, 8)
 GROUP_SIZES = (32, 64, 128)
@@ -65,7 +66,7 @@ def test_dq_gemm_matches_qmm_for_supported_bit_widths(bits: int, group_size: int
     x = mx.random.normal((16, in_features)).astype(mx.bfloat16)
     mx.eval(x)
     before = dq_gemm_engaged()
-    got = linear(x, quantized)
+    got = h3_linear(x, quantized)
     ref = _qmm(x, quantized)
     mx.eval(got, ref)
     assert dq_gemm_engaged() == before + 1
@@ -84,7 +85,7 @@ def test_dq_gemm_with_bias_and_batched_rows(monkeypatch: pytest.MonkeyPatch) ->
     x = mx.random.normal((2, 8, 128)).astype(mx.bfloat16)
     bias = mx.random.normal((64, )).astype(mx.bfloat16)
     mx.eval(x, bias)
-    got = linear(x, quantized, bias)
+    got = h3_linear(x, quantized, bias)
     ref = _qmm(x, quantized) + bias
     mx.eval(got, ref)
     assert dq_gemm_engaged() == 1
@@ -99,7 +100,7 @@ def test_dq_gemm_stays_on_qmm_below_threshold(monkeypatch: pytest.MonkeyPatch) -
     quantized = _try_quantize(64, 128, bits=6, group_size=64)
     x = mx.random.normal((32, 128)).astype(mx.bfloat16)
     mx.eval(x)
-    got = linear(x, quantized)
+    got = h3_linear(x, quantized)
     ref = _qmm(x, quantized)
     mx.eval(got, ref)
     assert dq_gemm_engaged() == 0
@@ -113,9 +114,23 @@ def test_dq_gemm_env_zero_disables_dispatch(monkeypatch: pytest.MonkeyPatch) ->
     quantized = _try_quantize(64, 128, bits=6, group_size=64)
     x = mx.random.normal((1024, 128)).astype(mx.bfloat16)
     mx.eval(x)
-    linear(x, quantized)
-    mx.eval(x)
+    got = h3_linear(x, quantized)
+    ref = _qmm(x, quantized)
+    mx.eval(got, ref)
+    assert dq_gemm_engaged() == 0
+    np.testing.assert_array_equal(np.asarray(got.astype(mx.float32)), np.asarray(ref.astype(mx.float32)))
+
+
+def test_shared_linear_stays_on_qmm_at_wide_m(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setenv("FASTVIDEO_MLX_DQ_GEMM", "1")
+    reset_dq_gemm_telemetry()
+    quantized = _try_quantize(64, 128, bits=6, group_size=64)
+    x = mx.random.normal((1024, 128)).astype(mx.bfloat16)
+    got = linear(x, quantized)
+    ref = _qmm(x, quantized)
+    mx.eval(got, ref)
     assert dq_gemm_engaged() == 0
+    np.testing.assert_array_equal(np.asarray(got.astype(mx.float32)), np.asarray(ref.astype(mx.float32)))
 
 
 def test_non_affine_weights_stay_on_quantized_matmul(monkeypatch: pytest.MonkeyPatch) -> None:
@@ -130,7 +145,7 @@ def test_non_affine_weights_stay_on_quantized_matmul(monkeypatch: pytest.MonkeyP
         pytest.skip(f"mxfp8 unsupported: {exc}")
     x = mx.random.normal((1024, 64)).astype(mx.bfloat16)
     mx.eval(x)
-    got = linear(x, quantized)
+    got = h3_linear(x, quantized)
     mx.eval(got)
     assert dq_gemm_engaged() == 0
     assert got.shape == (1024, 64)