| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """ROCm numerical-equivalence tests for the fused MTP depth-layer input kernel. |
| 4 | |
| 5 | ``embed_dual_rmsnorm_cat`` must match the unfused module sequence numerically: |
| 6 | each rmsnorm computes in fp32 and rounds to bf16 at the same points as the |
| 7 | vendored ``rmsnorm`` kernel (including the bf16 round-trip between the |
| 8 | chained backbone embed_norm and the depth embed_norm), and the fused row |
| 9 | gather matches ``F.embedding``. |
| 10 | """ |
| 11 | |
| 12 | from typing import cast |
| 13 | |
| 14 | import pytest |
| 15 | import torch |
| 16 | |
| 17 | from vllm.models.inkling.amd.ops.norm import ( |
| 18 | embed_dual_rmsnorm_cat, |
| 19 | embed_rmsnorm, |
| 20 | rmsnorm, |
| 21 | ) |
| 22 | |
| 23 | EPS = 1e-6 |
| 24 | VOCAB = 4096 |
| 25 | BF16_ATOL = torch.finfo(torch.bfloat16).eps |
| 26 | |
| 27 | |
| 28 | def _assert_bf16_close(actual: torch.Tensor, expected: torch.Tensor) -> None: |
| 29 | # Different Triton reduction schedules can move a small number of outputs |
| 30 | # by one bf16 epsilon while remaining numerically equivalent. |
| 31 | torch.testing.assert_close(actual, expected, rtol=0, atol=BF16_ATOL) |
| 32 | |
| 33 | |
| 34 | def _ref(hidden, w_h, w_e, emb, w_pre=None): |
| 35 | if w_pre is not None: |
| 36 | emb = rmsnorm(emb, w_pre, EPS) |
| 37 | return torch.cat([rmsnorm(hidden, w_h, EPS), rmsnorm(emb, w_e, EPS)], dim=-1) |
| 38 | |
| 39 | |
| 40 | @pytest.mark.parametrize("n", [1536, 6144]) |
| 41 | @pytest.mark.parametrize("t", [0, 1, 7, 256]) |
| 42 | @pytest.mark.parametrize("ids_dtype", [torch.int32, torch.int64]) |
| 43 | def test_embed_dual_rmsnorm_cat(n: int, t: int, ids_dtype: torch.dtype) -> None: |
| 44 | torch.manual_seed(0) |
| 45 | dev = "cuda" |
| 46 | table = (torch.randn(VOCAB, n, device=dev) * 0.3).to(torch.bfloat16) |
| 47 | w_h = torch.randn(n, device=dev).to(torch.bfloat16) |
| 48 | w_e = (1 + 0.01 * torch.randn(n, device=dev)).to(torch.bfloat16) |
| 49 | w_pre = torch.randn(n, device=dev).to(torch.bfloat16) |
| 50 | ids = torch.randint(0, VOCAB, (t,), device=dev, dtype=ids_dtype) |
| 51 | hidden = (torch.randn(t, n, device=dev) * 2).to(torch.bfloat16) |
| 52 | emb = table[ids.long()] |
| 53 | |
| 54 | # Fused gather + chained backbone pre-norm (the decode draft-step path). |
| 55 | out = embed_dual_rmsnorm_cat( |
| 56 | hidden, |
| 57 | w_h, |
| 58 | w_e, |
| 59 | EPS, |
| 60 | input_ids=ids, |
| 61 | embed_table=table, |
| 62 | pre_norm_weight=w_pre, |
| 63 | ) |
| 64 | assert out.shape == (t, 2 * n) |
| 65 | _assert_bf16_close(out, _ref(hidden, w_h, w_e, emb, w_pre)) |
| 66 | |
| 67 | # Precomputed embeds, no pre-norm (draft prefill with target-merged MM |
| 68 | # embeddings, already backbone-normed). |
| 69 | out = embed_dual_rmsnorm_cat(hidden, w_h, w_e, EPS, embeds=emb) |
| 70 | _assert_bf16_close(out, _ref(hidden, w_h, w_e, emb)) |
| 71 | |
| 72 | # Fused gather, no pre-norm (use_embed_norm=False). |
| 73 | out = embed_dual_rmsnorm_cat( |
| 74 | hidden, w_h, w_e, EPS, input_ids=ids, embed_table=table |
| 75 | ) |
| 76 | _assert_bf16_close(out, _ref(hidden, w_h, w_e, emb)) |
| 77 | |
| 78 | |
| 79 | @pytest.mark.parametrize("n", [1536, 6144]) |
| 80 | @pytest.mark.parametrize("t", [0, 1, 7, 256]) |
| 81 | @pytest.mark.parametrize("ids_dtype", [torch.int32, torch.int64]) |
| 82 | def test_embed_rmsnorm(n: int, t: int, ids_dtype: torch.dtype) -> None: |
| 83 | torch.manual_seed(0) |
| 84 | dev = "cuda" |
| 85 | table = (torch.randn(VOCAB, n, device=dev) * 0.3).to(torch.bfloat16) |
| 86 | w = torch.randn(n, device=dev).to(torch.bfloat16) |
| 87 | ids = torch.randint(0, VOCAB, (t,), device=dev, dtype=ids_dtype) |
| 88 | ref_emb = table[ids.long()] |
| 89 | |
| 90 | # Gather + embed_norm (base model / MTP prefill embed path). |
| 91 | out = cast(torch.Tensor, embed_rmsnorm(ids, table, w, EPS)) |
| 92 | assert out.shape == (t, n) |
| 93 | _assert_bf16_close(out, rmsnorm(ref_emb, w, EPS) if t else ref_emb) |
| 94 | |
| 95 | # Pure gather (use_embed_norm=False / replicated module forward). |
| 96 | out = cast(torch.Tensor, embed_rmsnorm(ids, table, None, EPS)) |
| 97 | assert torch.equal(out, ref_emb) |
| 98 | |
| 99 | # Chained first-layer attn_norm (the target text-path forward): one launch |
| 100 | # emits both the residual and layer 0's normed attention input. |
| 101 | w_chain = (1 + 0.05 * torch.randn(n, device=dev)).to(torch.bfloat16) |
| 102 | res, attn_in = cast( |
| 103 | tuple[torch.Tensor, torch.Tensor], |
| 104 | embed_rmsnorm(ids, table, w, EPS, chain_weight=w_chain), |
| 105 | ) |
| 106 | ref_res = rmsnorm(ref_emb, w, EPS) if t else ref_emb |
| 107 | _assert_bf16_close(res, ref_res) |
| 108 | _assert_bf16_close(attn_in, rmsnorm(ref_res, w_chain, EPS) if t else ref_res) |
| 109 | |
| 110 | # Chained without embed_norm (use_embed_norm=False). |
| 111 | res, attn_in = cast( |
| 112 | tuple[torch.Tensor, torch.Tensor], |
| 113 | embed_rmsnorm(ids, table, None, EPS, chain_weight=w_chain), |
| 114 | ) |
| 115 | assert torch.equal(res, ref_emb) |
| 116 | _assert_bf16_close(attn_in, rmsnorm(ref_emb, w_chain, EPS) if t else ref_emb) |