| 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)) |
| ) |
|
|
