| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Bit-exactness tests for the fused MTP depth-layer input kernel. |
| 4 | |
| 5 | ``embed_dual_rmsnorm_cat`` must match the unfused module sequence exactly: |
| 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 | import pytest |
| 13 | import torch |
| 14 | |
| 15 | from vllm.platforms import current_platform |
| 16 | |
| 17 | if not current_platform.is_cuda(): |
| 18 | pytest.skip("requires CUDA", allow_module_level=True) |
| 19 | |
| 20 | from vllm.models.inkling.nvidia.ops.norm import ( |
| 21 | embed_dual_rmsnorm_cat, |
| 22 | embed_rmsnorm, |
| 23 | rmsnorm, |
| 24 | ) |
| 25 | |
| 26 | EPS = 1e-6 |
| 27 | VOCAB = 4096 |
| 28 | |
| 29 | |
| 30 | def _ref(hidden, w_h, w_e, emb, w_pre=None): |
| 31 | if w_pre is not None: |
| 32 | emb = rmsnorm(emb, w_pre, EPS) |
| 33 | return torch.cat([rmsnorm(hidden, w_h, EPS), rmsnorm(emb, w_e, EPS)], dim=-1) |
| 34 | |
| 35 | |
| 36 | @pytest.mark.parametrize("n", [1536, 6144]) |
| 37 | @pytest.mark.parametrize("t", [0, 1, 7, 256]) |
| 38 | @pytest.mark.parametrize("ids_dtype", [torch.int32, torch.int64]) |
| 39 | def test_embed_dual_rmsnorm_cat(n: int, t: int, ids_dtype: torch.dtype) -> None: |
| 40 | torch.manual_seed(0) |
| 41 | dev = "cuda" |
| 42 | table = (torch.randn(VOCAB, n, device=dev) * 0.3).to(torch.bfloat16) |
| 43 | w_h = torch.randn(n, device=dev).to(torch.bfloat16) |
| 44 | w_e = (1 + 0.01 * torch.randn(n, device=dev)).to(torch.bfloat16) |
| 45 | w_pre = torch.randn(n, device=dev).to(torch.bfloat16) |
| 46 | ids = torch.randint(0, VOCAB, (t,), device=dev, dtype=ids_dtype) |
| 47 | hidden = (torch.randn(t, n, device=dev) * 2).to(torch.bfloat16) |
| 48 | emb = table[ids.long()] |
| 49 | |
| 50 | # Fused gather + chained backbone pre-norm (the decode draft-step path). |
| 51 | out = embed_dual_rmsnorm_cat( |
| 52 | hidden, |
| 53 | w_h, |
| 54 | w_e, |
| 55 | EPS, |
| 56 | input_ids=ids, |
| 57 | embed_table=table, |
| 58 | pre_norm_weight=w_pre, |
| 59 | ) |
| 60 | assert out.shape == (t, 2 * n) |
| 61 | assert torch.equal(out, _ref(hidden, w_h, w_e, emb, w_pre)) |
| 62 | |
| 63 | # Precomputed embeds, no pre-norm (draft prefill with target-merged MM |
| 64 | # embeddings, already backbone-normed). |
| 65 | out = embed_dual_rmsnorm_cat(hidden, w_h, w_e, EPS, embeds=emb) |
| 66 | assert torch.equal(out, _ref(hidden, w_h, w_e, emb)) |
| 67 | |
| 68 | # Fused gather, no pre-norm (use_embed_norm=False). |
| 69 | out = embed_dual_rmsnorm_cat( |
| 70 | hidden, w_h, w_e, EPS, input_ids=ids, embed_table=table |
| 71 | ) |
| 72 | assert torch.equal(out, _ref(hidden, w_h, w_e, emb)) |
| 73 | |
| 74 | |
| 75 | @pytest.mark.parametrize("n", [1536, 6144]) |
| 76 | @pytest.mark.parametrize("t", [0, 1, 7, 256]) |
| 77 | @pytest.mark.parametrize("ids_dtype", [torch.int32, torch.int64]) |
| 78 | def test_embed_rmsnorm(n: int, t: int, ids_dtype: torch.dtype) -> None: |
| 79 | torch.manual_seed(0) |
| 80 | dev = "cuda" |
| 81 | table = (torch.randn(VOCAB, n, device=dev) * 0.3).to(torch.bfloat16) |
| 82 | w = torch.randn(n, device=dev).to(torch.bfloat16) |
| 83 | ids = torch.randint(0, VOCAB, (t,), device=dev, dtype=ids_dtype) |
| 84 | ref_emb = table[ids.long()] |
| 85 | |
| 86 | # Gather + embed_norm (base model / MTP prefill embed path). |
| 87 | out = embed_rmsnorm(ids, table, w, EPS) |
| 88 | assert out.shape == (t, n) |
| 89 | assert torch.equal(out, rmsnorm(ref_emb, w, EPS) if t else ref_emb) |
| 90 | |
| 91 | # Pure gather (use_embed_norm=False / replicated module forward). |
| 92 | out = embed_rmsnorm(ids, table, None, EPS) |
| 93 | assert torch.equal(out, ref_emb) |
| 94 | |
| 95 | # Chained first-layer attn_norm (the target text-path forward): one launch |
| 96 | # emits both the residual and layer 0's normed attention input. |
| 97 | w_chain = (1 + 0.05 * torch.randn(n, device=dev)).to(torch.bfloat16) |
| 98 | res, attn_in = embed_rmsnorm(ids, table, w, EPS, chain_weight=w_chain) |
| 99 | ref_res = rmsnorm(ref_emb, w, EPS) if t else ref_emb |
| 100 | assert torch.equal(res, ref_res) |
| 101 | assert torch.equal(attn_in, rmsnorm(ref_res, w_chain, EPS) if t else ref_res) |
| 102 | |
| 103 | # Chained without embed_norm (use_embed_norm=False). |
| 104 | res, attn_in = embed_rmsnorm(ids, table, None, EPS, chain_weight=w_chain) |
| 105 | assert torch.equal(res, ref_emb) |
| 106 | assert torch.equal(attn_in, rmsnorm(ref_emb, w_chain, EPS) if t else ref_emb) |