| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import pytest |
| 4 | import torch |
| 5 | |
| 6 | from vllm.model_executor.models.gemma4 import ( |
| 7 | gemma4_fused_routing_kernel_triton, |
| 8 | gemma4_routing_function_torch, |
| 9 | ) |
| 10 | |
| 11 | |
| 12 | def sort_by_id(w, ids): |
| 13 | order = ids.argsort(dim=-1) |
| 14 | return w.gather(1, order), ids.gather(1, order) |
| 15 | |
| 16 | |
| 17 | # Gemma4 MoE Model has context length of 250K |
| 18 | # the minus 1 is to ensure that edge cases are tested |
| 19 | @pytest.mark.parametrize("num_tokens", [1, 2, 2048, 250000]) |
| 20 | @pytest.mark.parametrize("num_experts", [128]) # gemma4 moe experts |
| 21 | @pytest.mark.parametrize("topk", [8]) # gemma4 topk |
| 22 | @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.half, torch.float32]) |
| 23 | def test_gemma4_routing_kernel_triton( |
| 24 | num_tokens: int, |
| 25 | num_experts: int, |
| 26 | topk: int, |
| 27 | dtype: torch.dtype, |
| 28 | ): |
| 29 | torch.manual_seed(0) |
| 30 | |
| 31 | gating = torch.randn(num_tokens, num_experts, dtype=dtype, device="cuda") |
| 32 | scales = torch.rand(num_experts, dtype=torch.float32, device="cuda") |
| 33 | |
| 34 | ref_w, ref_ids = gemma4_routing_function_torch(gating, topk, scales) |
| 35 | tri_w, tri_ids = gemma4_fused_routing_kernel_triton(gating, topk, scales) |
| 36 | |
| 37 | # Sort by expert id — to remove tie-breaking differences |
| 38 | ref_ws, ref_is = sort_by_id(ref_w, ref_ids) |
| 39 | tri_ws, tri_is = sort_by_id(tri_w, tri_ids) |
| 40 | |
| 41 | ids_match = (ref_is == tri_is).all().item() |
| 42 | weights_match = torch.allclose(ref_ws, tri_ws, atol=1e-2, rtol=1e-2) |
| 43 | all_match = ids_match and weights_match |
| 44 | max_err = (ref_ws - tri_ws).abs().max().item() |
| 45 | print( |
| 46 | f"T={num_tokens:5d} E={num_experts:4d} K={topk} " |
| 47 | f"{str(dtype).split('.')[-1]:7s} ids={ids_match} max_Δweight={max_err:.2e}" |
| 48 | ) |
| 49 | if not all_match: |
| 50 | bad = (ref_is != tri_is).any(dim=-1).nonzero(as_tuple=True)[0] |
| 51 | if len(bad): |
| 52 | r = bad[0].item() |
| 53 | print( |
| 54 | f" first bad row {r}: ref_ids={ref_ids[r].tolist()} " |
| 55 | f"tri_ids={tri_ids[r].tolist()}" |
| 56 | ) |
| 57 | assert all_match |