| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Unit tests for Quark OCP MXFP4 weights loaded into AITER MoE layouts.""" |
| 4 | |
| 5 | from types import SimpleNamespace |
| 6 | |
| 7 | import torch |
| 8 | from torch import nn |
| 9 | |
| 10 | from vllm.models.inkling.amd.moe import InklingMoE |
| 11 | |
| 12 | |
| 13 | def _fake_moe() -> tuple[InklingMoE, SimpleNamespace]: |
| 14 | moe = InklingMoE.__new__(InklingMoE) |
| 15 | nn.Module.__init__(moe) |
| 16 | moe.n_routed_experts = 4 |
| 17 | routed = SimpleNamespace( |
| 18 | moe_config=SimpleNamespace( |
| 19 | moe_parallel_config=SimpleNamespace(tp_rank=1, tp_size=2) |
| 20 | ), |
| 21 | # Logical intermediate per rank is 3; AITER pads it to 4. |
| 22 | w13_weight=nn.Parameter(torch.zeros(4, 8, 2, dtype=torch.uint8), False), |
| 23 | w2_weight=nn.Parameter(torch.zeros(4, 5, 3, dtype=torch.uint8), False), |
| 24 | w13_weight_scale=nn.Parameter(torch.ones(4, 8, 2, dtype=torch.uint8), False), |
| 25 | w2_weight_scale=nn.Parameter(torch.ones(4, 5, 3, dtype=torch.uint8), False), |
| 26 | ) |
| 27 | moe.experts = SimpleNamespace(routed_experts=routed) |
| 28 | moe._local_expert_slots = ( # type: ignore[method-assign] |
| 29 | lambda: dict(enumerate(range(4))) |
| 30 | ) |
| 31 | return moe, routed |
| 32 | |
| 33 | |
| 34 | def test_loads_logical_tp_shards_and_preserves_aiter_padding(): |
| 35 | moe, routed = _fake_moe() |
| 36 | |
| 37 | w13 = torch.arange(4 * 12 * 2, dtype=torch.uint8).view(4, 12, 2) |
| 38 | w2 = torch.arange(4 * 5 * 4, dtype=torch.uint8).view(4, 5, 4) |
| 39 | moe.load_expert_weight("experts.w13_weight", w13) |
| 40 | moe.load_expert_weight("experts.w2_weight", w2) |
| 41 | |
| 42 | # TP rank 1 consumes the second six interleaved gate/up rows. AITER keeps |
| 43 | # one padding row in each destination half and one padding w2 column. |
| 44 | torch.testing.assert_close(routed.w13_weight[:, :3], w13[:, 6:12:2]) |
| 45 | torch.testing.assert_close(routed.w13_weight[:, 4:7], w13[:, 7:12:2]) |
| 46 | assert torch.count_nonzero(routed.w13_weight[:, 3]) == 0 |
| 47 | assert torch.count_nonzero(routed.w13_weight[:, 7]) == 0 |
| 48 | torch.testing.assert_close(routed.w2_weight[:, :, :2], w2[:, :, 2:4]) |
| 49 | assert torch.count_nonzero(routed.w2_weight[:, :, 2]) == 0 |
| 50 | |
| 51 | |
| 52 | def test_unflattens_quark_scales_and_preserves_scale_padding(): |
| 53 | moe, routed = _fake_moe() |
| 54 | |
| 55 | flat_w13_scale = torch.arange(4 * 12 * 2, dtype=torch.uint8).view(4 * 12, 2) |
| 56 | flat_w2_scale = torch.arange(4 * 5 * 4, dtype=torch.uint8).view(4 * 5, 4) |
| 57 | moe.load_expert_weight("experts.w13_weight_scale", flat_w13_scale) |
| 58 | moe.load_expert_weight("experts.w2_weight_scale", flat_w2_scale) |
| 59 | |
| 60 | w13_scale = flat_w13_scale.view(4, 12, 2) |
| 61 | w2_scale = flat_w2_scale.view(4, 5, 4) |
| 62 | torch.testing.assert_close(routed.w13_weight_scale[:, :3], w13_scale[:, 6:12:2]) |
| 63 | torch.testing.assert_close(routed.w13_weight_scale[:, 4:7], w13_scale[:, 7:12:2]) |
| 64 | assert torch.all(routed.w13_weight_scale[:, 3] == 1) |
| 65 | assert torch.all(routed.w13_weight_scale[:, 7] == 1) |
| 66 | torch.testing.assert_close(routed.w2_weight_scale[:, :, :2], w2_scale[:, :, 2:4]) |
| 67 | assert torch.all(routed.w2_weight_scale[:, :, 2] == 1) |