| import mlx.core as mx |
| import numpy as np |
|
|
| from h3_mlx.config import DiTConfig |
| from h3_mlx.dit import DiTBlock |
| from h3_mlx.weights import WeightStore |
|
|
|
|
| def test_packed_block_cache_and_compiled_bf16_boundaries(tmp_path): |
| mx.random.seed(729) |
| c = DiTConfig( |
| hidden_size=32, num_attention_heads=2, attention_head_dim=16, ffn_hidden_size=64 |
| ) |
| tensors = {} |
| for name, shape in { |
| "norm1": (32,), |
| "norm2": (32,), |
| "attn.q_norm": (16,), |
| "attn.k_norm": (16,), |
| "attn.qkv_proj": (96, 32), |
| "attn.out_proj": (32, 32), |
| "mlp.fc1": (128, 32), |
| "mlp.fc2": (32, 64), |
| }.items(): |
| tensors[f"blocks.0.{name}.weight"] = (mx.random.normal(shape) * 0.02).astype( |
| mx.bfloat16 |
| ) |
| if len(shape) == 2: |
| tensors[f"blocks.0.{name}.bias"] = ( |
| mx.random.normal((shape[0],)) * 0.01 |
| ).astype(mx.bfloat16) |
| original = tmp_path / "original" |
| original.mkdir() |
| mx.save_safetensors(str(original / "model.safetensors"), tensors) |
| store = WeightStore(original) |
| native = DiTBlock(store, "blocks.0", c) |
| packed = DiTBlock(store, "blocks.0", c, packed_cache=tmp_path / "packed") |
| restored = DiTBlock(store, "blocks.0", c, packed_cache=tmp_path / "packed") |
| x = mx.random.normal((17, 32)).astype(mx.bfloat16) |
| params = tuple( |
| (mx.random.normal((3, 32)) * 0.1).astype(mx.bfloat16) for _ in range(6) |
| ) |
| ids = mx.arange(17) % 3 |
| expected = native(x, params=params, indices=ids, mlp_chunk=0) |
| state = tuple( |
| p.weight.state() |
| for p in (restored.qkv, restored.out, restored.fc1, restored.fc2) |
| ) |
|
|
| def forward(x, params, ids, state): |
| return restored( |
| x, params=params, indices=ids, mlp_chunk=0, projection_weights=state |
| ) |
|
|
| compiled = mx.compile(forward) |
| for actual in ( |
| packed(x, params=params, indices=ids, mlp_chunk=0), |
| restored(x, params=params, indices=ids, mlp_chunk=0), |
| compiled(x, params, ids, state), |
| ): |
| np.testing.assert_array_equal( |
| np.asarray(actual.view(mx.uint16)), np.asarray(expected.view(mx.uint16)) |
| ) |
|
|
