File size: 3,039 Bytes
3da88e3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 | import mlx.core as mx
import numpy as np
import pytest
from h3_mlx.ops import modulate, rms_norm, rotate
from h3_mlx.tuning import (
TunedFC2,
_gemm,
norm_modulate_exact,
qk_norm_rope_exact,
supported_device,
)
from h3_mlx.weights import Linear
@pytest.mark.parametrize("magnitude", [0.001, 1.0, 10000.0])
@pytest.mark.parametrize("rope_dtype", [mx.bfloat16, mx.float32])
@pytest.mark.skipif(
not supported_device(), reason="Pinned M5 Pro/MLX kernel validation"
)
def test_qk_fusion_preserves_native_reduction_and_rope_bits(magnitude, rope_dtype):
mx.random.seed(8192)
qkv = (mx.random.normal((37, 3, 3, 128)) * magnitude).astype(mx.bfloat16)
qw = mx.random.normal((128,)).astype(mx.bfloat16)
kw = mx.random.normal((128,)).astype(mx.bfloat16)
angles = mx.random.normal((1, 37, 1, 48))
angles = mx.concatenate((angles, angles), axis=-1)
cosine = mx.cos(angles).astype(rope_dtype)
sine = mx.sin(angles).astype(rope_dtype)
reference = tuple(
rotate(rms_norm(qkv[None, :, i], w), cosine, sine)
for i, w in ((0, qw), (1, kw))
)
result = qk_norm_rope_exact(qkv, qw, kw, cosine, sine)
for expected, actual in zip(reference, result, strict=True):
np.testing.assert_array_equal(
np.asarray(actual.view(mx.uint16)), np.asarray(expected.view(mx.uint16))
)
@pytest.mark.skipif(not supported_device(), reason="Metal 4 M5 kernel validation")
def test_mpp_fp32_bias_epilogue_and_partial_m_tile():
mx.random.seed(4352)
x = (mx.random.normal((71, 512)) * 32).astype(mx.bfloat16)
weight = (mx.random.normal((256, 512)) * 0.02).astype(mx.bfloat16)
bias = mx.random.normal((256,)).astype(mx.bfloat16)
expected = mx.addmm(bias, x, weight.T)
actual = _gemm(x, weight, bias, 64, 128, 128, 4, 8)
np.testing.assert_array_equal(
np.asarray(actual.view(mx.uint16)), np.asarray(expected.view(mx.uint16))
)
def test_tuned_fc2_keeps_native_fallback_for_unsupported_shapes():
mx.random.seed(71)
layer = Linear(mx.random.normal((32, 64)).astype(mx.bfloat16))
x = mx.random.normal((17, 64)).astype(mx.bfloat16)
tuned = TunedFC2(layer)
np.testing.assert_array_equal(
np.asarray(tuned(x).view(mx.uint16)), np.asarray(layer(x).view(mx.uint16))
)
@pytest.mark.skipif(not supported_device(), reason="Pinned M5 Pro reduction layout")
def test_full_width_norm_fusion_matches_native_loop_and_bf16_boundaries():
mx.random.seed(5376)
x = (mx.random.normal((13, 5376)) * 256).astype(mx.bfloat16)
weight = mx.random.normal((5376,)).astype(mx.bfloat16)
shift = mx.random.normal((6, 5376)).astype(mx.bfloat16)
scale = (mx.random.normal((6, 5376)) * 0.1).astype(mx.bfloat16)
indices = mx.arange(13, dtype=mx.int32) % 6
expected = modulate(x, weight, shift, scale, indices)
actual = norm_modulate_exact(x, weight, shift, scale, indices)
np.testing.assert_array_equal(
np.asarray(actual.view(mx.uint16)), np.asarray(expected.view(mx.uint16))
)
|