forked from Karylab-cklius/vllm
+1
fb5ec0dc9e
fb5ec0dc9e
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai> Co-authored-by: Bugen Zhao <i@bugenzhao.com> Co-authored-by: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> Co-authored-by: Isotr0py <Isotr0py@outlook.com> Co-authored-by: Isotr0py <mozf@inferact.ai> Co-authored-by: Jee Jee Li <jeejeelee@inferact.ai> Co-authored-by: Roger Wang <hey@rogerw.io> Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: OpenAI Codex <codex@openai.com>
107 lines
4.0 KiB
Python
107 lines
4.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Bit-exactness tests for the fused MTP depth-layer input kernel.
|
|
|
|
``embed_dual_rmsnorm_cat`` must match the unfused module sequence exactly:
|
|
each rmsnorm computes in fp32 and rounds to bf16 at the same points as the
|
|
vendored ``rmsnorm`` kernel (including the bf16 round-trip between the
|
|
chained backbone embed_norm and the depth embed_norm), and the fused row
|
|
gather matches ``F.embedding``.
|
|
"""
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm.platforms import current_platform
|
|
|
|
if not current_platform.is_cuda():
|
|
pytest.skip("requires CUDA", allow_module_level=True)
|
|
|
|
from vllm.models.inkling.nvidia.ops.norm import (
|
|
embed_dual_rmsnorm_cat,
|
|
embed_rmsnorm,
|
|
rmsnorm,
|
|
)
|
|
|
|
EPS = 1e-6
|
|
VOCAB = 4096
|
|
|
|
|
|
def _ref(hidden, w_h, w_e, emb, w_pre=None):
|
|
if w_pre is not None:
|
|
emb = rmsnorm(emb, w_pre, EPS)
|
|
return torch.cat([rmsnorm(hidden, w_h, EPS), rmsnorm(emb, w_e, EPS)], dim=-1)
|
|
|
|
|
|
@pytest.mark.parametrize("n", [1536, 6144])
|
|
@pytest.mark.parametrize("t", [0, 1, 7, 256])
|
|
@pytest.mark.parametrize("ids_dtype", [torch.int32, torch.int64])
|
|
def test_embed_dual_rmsnorm_cat(n: int, t: int, ids_dtype: torch.dtype) -> None:
|
|
torch.manual_seed(0)
|
|
dev = "cuda"
|
|
table = (torch.randn(VOCAB, n, device=dev) * 0.3).to(torch.bfloat16)
|
|
w_h = torch.randn(n, device=dev).to(torch.bfloat16)
|
|
w_e = (1 + 0.01 * torch.randn(n, device=dev)).to(torch.bfloat16)
|
|
w_pre = torch.randn(n, device=dev).to(torch.bfloat16)
|
|
ids = torch.randint(0, VOCAB, (t,), device=dev, dtype=ids_dtype)
|
|
hidden = (torch.randn(t, n, device=dev) * 2).to(torch.bfloat16)
|
|
emb = table[ids.long()]
|
|
|
|
# Fused gather + chained backbone pre-norm (the decode draft-step path).
|
|
out = embed_dual_rmsnorm_cat(
|
|
hidden,
|
|
w_h,
|
|
w_e,
|
|
EPS,
|
|
input_ids=ids,
|
|
embed_table=table,
|
|
pre_norm_weight=w_pre,
|
|
)
|
|
assert out.shape == (t, 2 * n)
|
|
assert torch.equal(out, _ref(hidden, w_h, w_e, emb, w_pre))
|
|
|
|
# Precomputed embeds, no pre-norm (draft prefill with target-merged MM
|
|
# embeddings, already backbone-normed).
|
|
out = embed_dual_rmsnorm_cat(hidden, w_h, w_e, EPS, embeds=emb)
|
|
assert torch.equal(out, _ref(hidden, w_h, w_e, emb))
|
|
|
|
# Fused gather, no pre-norm (use_embed_norm=False).
|
|
out = embed_dual_rmsnorm_cat(
|
|
hidden, w_h, w_e, EPS, input_ids=ids, embed_table=table
|
|
)
|
|
assert torch.equal(out, _ref(hidden, w_h, w_e, emb))
|
|
|
|
|
|
@pytest.mark.parametrize("n", [1536, 6144])
|
|
@pytest.mark.parametrize("t", [0, 1, 7, 256])
|
|
@pytest.mark.parametrize("ids_dtype", [torch.int32, torch.int64])
|
|
def test_embed_rmsnorm(n: int, t: int, ids_dtype: torch.dtype) -> None:
|
|
torch.manual_seed(0)
|
|
dev = "cuda"
|
|
table = (torch.randn(VOCAB, n, device=dev) * 0.3).to(torch.bfloat16)
|
|
w = torch.randn(n, device=dev).to(torch.bfloat16)
|
|
ids = torch.randint(0, VOCAB, (t,), device=dev, dtype=ids_dtype)
|
|
ref_emb = table[ids.long()]
|
|
|
|
# Gather + embed_norm (base model / MTP prefill embed path).
|
|
out = embed_rmsnorm(ids, table, w, EPS)
|
|
assert out.shape == (t, n)
|
|
assert torch.equal(out, rmsnorm(ref_emb, w, EPS) if t else ref_emb)
|
|
|
|
# Pure gather (use_embed_norm=False / replicated module forward).
|
|
out = embed_rmsnorm(ids, table, None, EPS)
|
|
assert torch.equal(out, ref_emb)
|
|
|
|
# Chained first-layer attn_norm (the target text-path forward): one launch
|
|
# emits both the residual and layer 0's normed attention input.
|
|
w_chain = (1 + 0.05 * torch.randn(n, device=dev)).to(torch.bfloat16)
|
|
res, attn_in = embed_rmsnorm(ids, table, w, EPS, chain_weight=w_chain)
|
|
ref_res = rmsnorm(ref_emb, w, EPS) if t else ref_emb
|
|
assert torch.equal(res, ref_res)
|
|
assert torch.equal(attn_in, rmsnorm(ref_res, w_chain, EPS) if t else ref_res)
|
|
|
|
# Chained without embed_norm (use_embed_norm=False).
|
|
res, attn_in = embed_rmsnorm(ids, table, None, EPS, chain_weight=w_chain)
|
|
assert torch.equal(res, ref_emb)
|
|
assert torch.equal(attn_in, rmsnorm(ref_emb, w_chain, EPS) if t else ref_emb)
|