iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Replicated embedding and norm fusion for DSV3 flat model (#48484)

Browse Source 
Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>
This commit is contained in:
Jee Jee Li
2026-08-18 14:53:21 +08:00
committed by GitHub
parent 5fa8ca971a
commit 2687fec6ef
5 changed files with 385 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/kernels/core/test_fused_embed_norm.py
+75
View File
@@ -0,0 +1,75 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for the replicated-embedding fused gather/norm kernels
(``vllm.model_executor.layers.fused_embed_norm``).
The guarantee: enabling ``VLLM_REPLICATE_EMBED`` (replicated table + fused
kernels) must not change model outputs. The gathered residual is bit-exact and
the fused norms match the unfused reference.
"""
import pytest
import torch
from vllm.model_executor.layers.fused_embed_norm import (
fused_embed_eh_norm,
fused_embed_norm,
)
# The model-local (untouched) eh-norm the replicate path must match.
from vllm.models.deepseek_v32.common.kernels import fused_eh_norm
from vllm.platforms import current_platform
from vllm.utils.torch_utils import set_random_seed
DTYPE = torch.bfloat16
VOCAB, HIDDEN, NUM_TOKENS, EPS = 8192, 4096, 129, 1e-6
requires_cuda = pytest.mark.skipif(
not current_platform.is_cuda_alike(),
reason="fused embed/norm Triton kernels require a CUDA/ROCm device",
)
def _rmsnorm(x: torch.Tensor, w: torch.Tensor, eps: float) -> torch.Tensor:
# Full-precision (fp32) reference RMSNorm.
var = x.float().pow(2).mean(dim=-1, keepdim=True)
return x.float() * torch.rsqrt(var + eps) * w.float()
@requires_cuda
@torch.inference_mode()
def test_fused_embed_norm_matches_reference():
"""Main-model fusion: the residual is the exact gather and the second output
is a correct RMSNorm. The norm matches a full-precision reference to ~2 bf16
ulp (rtol 1e-2) -- that gap is bf16 rounding, not the kernel."""
set_random_seed(13)
table = torch.randn(VOCAB, HIDDEN, dtype=DTYPE, device="cuda")
ids = torch.randint(0, VOCAB, (NUM_TOKENS,), dtype=torch.int32, device="cuda")
weight = torch.empty(HIDDEN, dtype=DTYPE, device="cuda").normal_(1.0, 0.1)
residual, normed = fused_embed_norm(ids, table, chain_weight=weight, eps=EPS)
embeds = table[ids.long()]
torch.testing.assert_close(residual, embeds, atol=0.0, rtol=0.0)
torch.testing.assert_close(
normed.float(), _rmsnorm(embeds, weight, EPS), atol=1e-3, rtol=1e-2
)
@requires_cuda
@torch.inference_mode()
def test_fused_embed_eh_norm_matches_reference():
"""MTP fusion (folded gather) is bit-exact vs gathering the embeds and
feeding the untouched model-local ``fused_eh_norm``."""
set_random_seed(13)
table = torch.randn(VOCAB, HIDDEN, dtype=DTYPE, device="cuda")
ids = torch.randint(0, VOCAB, (NUM_TOKENS,), dtype=torch.int32, device="cuda")
prev = torch.randn(NUM_TOKENS, HIDDEN, dtype=DTYPE, device="cuda")
enorm_w = torch.randn(HIDDEN, dtype=DTYPE, device="cuda")
hnorm_w = torch.randn(HIDDEN, dtype=DTYPE, device="cuda")
positions = torch.arange(NUM_TOKENS, device="cuda") # includes pos 0
fused = fused_embed_eh_norm(positions, ids, table, prev, enorm_w, hnorm_w, EPS)
ref = fused_eh_norm(positions, table[ids.long()], prev, enorm_w, hnorm_w, EPS)
torch.testing.assert_close(fused, ref, atol=0.0, rtol=0.0)
vllm/envs.py
+4
View File
@@ -156,6 +156,7 @@ if TYPE_CHECKING:
VLLM_ENABLE_V1_MULTIPROCESSING: bool = True
VLLM_LOG_BATCHSIZE_INTERVAL: float = -1
VLLM_DISABLE_COMPILE_CACHE: bool = False
VLLM_REPLICATE_EMBED: bool = False
VLLM_USE_LAYERNAME: bool = True
Q_SCALE_CONSTANT: int = 200
K_SCALE_CONSTANT: int = 200
@@ -620,6 +621,9 @@ environment_variables: dict[str, Callable[[], Any]] = {
# Enable batch-invariant mode: deterministic results regardless of
# batch composition. Requires NVIDIA GPU with compute capability >= 9.0.
"VLLM_BATCH_INVARIANT": lambda: bool(int(os.getenv("VLLM_BATCH_INVARIANT", "0"))),
"VLLM_REPLICATE_EMBED": lambda: (
os.getenv("VLLM_REPLICATE_EMBED", "0").strip().lower() in ("1", "true")
),
# Use tensor descriptors for Q/K/V loads and output stores in the
# Triton unified-attention kernel. Enables HW 2D block reads on
# Intel XPU; the non-TD branch is dead-code-eliminated at Triton
vllm/model_executor/layers/fused_embed_norm.py
+231
View File
@@ -0,0 +1,231 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Replicated input embedding + its fused gather/norm kernels.
Groups the ``VLLM_REPLICATE_EMBED`` path in one place: the embedding factory,
the predicate that says whether the fusions apply, and the two Triton fusions
the full on-rank table unlocks --
* ``fused_embed_norm``: gather + a chained RMSNorm (e.g. the first decoder
layer's ``input_layernorm``), and
* ``fused_embed_eh_norm``: gather + pos-0 zeroing + enorm/hnorm + cat, the
embed/previous-hidden input norm for a speculative (MTP/eagle) depth layer
(the replicated-table analogue of the model-local ``fused_eh_norm``, which
takes precomputed embeds).
Self-contained (no model-local imports) so it can live under ``layers/``.
"""
import torch
import vllm.envs as envs
from vllm.distributed import get_tensor_model_parallel_world_size
from vllm.model_executor.layers.vocab_parallel_embedding import (
UnquantizedEmbeddingMethod,
VocabParallelEmbedding,
)
from vllm.triton_utils import tl, triton
@triton.jit
def _rms_norm(x, w, eps, HIDDEN_SIZE: tl.constexpr):
x = x.to(tl.float32)
mean_sq = tl.sum(x * x, axis=0) / HIDDEN_SIZE
rrms = tl.rsqrt(mean_sq + eps)
w = w.to(tl.float32)
return (x * rrms) * w
def make_input_embedding(
num_embeddings: int,
embedding_dim: int,
*,
params_dtype: torch.dtype | None = None,
quant_config=None,
prefix: str = "",
tie_word_embeddings: bool = False,
) -> VocabParallelEmbedding:
"""Input token embedding with an optional replicated escape hatch.
``VLLM_REPLICATE_EMBED=1`` builds the embedding with ``disable_tp``: the full
table lives on every rank and the lookup is a local gather with no mask and
no all-reduce, which unlocks the fused gather+norm path. The cost is a full
table per rank at TP>1 (no extra memory at TP=1, where vocab-parallel is
already unsharded). A replicated, unsharded table cannot be tied to a
vocab-parallel ``ParallelLMHead``, so tied word embeddings are rejected at
TP>1 (at TP=1 ``disable_tp`` is a no-op and tying still works).
"""
disable_tp = envs.VLLM_REPLICATE_EMBED
if disable_tp and tie_word_embeddings:
assert get_tensor_model_parallel_world_size() == 1, (
"VLLM_REPLICATE_EMBED is unsupported with tied word embeddings "
"(the replicated table cannot tie to a vocab-parallel lm_head)"
)
return VocabParallelEmbedding(
num_embeddings,
embedding_dim,
params_dtype=params_dtype,
quant_config=quant_config,
prefix=prefix,
disable_tp=disable_tp,
)
def has_full_vocab_on_rank(embedding: torch.nn.Module) -> bool:
"""Whether ``embedding.weight`` is the whole vocab as a plain [V, H] table.
The fused gather kernels index the table directly, so they need every row
on-rank (``disable_tp``, or any TP=1 run) and an unquantized weight.
"""
return getattr(embedding, "tp_size", 0) == 1 and isinstance(
getattr(embedding, "quant_method", None), UnquantizedEmbeddingMethod
)
@triton.jit
def _fused_embed_norm_kernel(
ids_ptr, # [T] token ids
table_ptr, # [V, H] embedding table (full vocab, replicated on-rank)
table_stride_0,
out_ptr, # [T, H] gathered embedding (the residual stream)
normed_ptr, # [T, H] rmsnorm(out, chain_w) (HAS_NORM only)
chain_w_ptr, # [H] next norm weight (HAS_NORM only)
eps,
H: tl.constexpr,
BLOCK: tl.constexpr,
HAS_NORM: tl.constexpr,
):
tok = tl.program_id(0).to(tl.int64)
off = tl.arange(0, BLOCK)
mask = off < H
row = tl.load(ids_ptr + tok).to(tl.int64)
x = tl.load(table_ptr + row * table_stride_0 + off, mask=mask, other=0.0)
tl.store(out_ptr + tok * H + off, x, mask=mask)
if HAS_NORM:
w = tl.load(chain_w_ptr + off, mask=mask)
y = _rms_norm(x, w, eps, H).to(normed_ptr.dtype.element_ty)
tl.store(normed_ptr + tok * H + off, y, mask=mask)
# Base model fusion
def fused_embed_norm(
input_ids: torch.Tensor,
embed_table: torch.Tensor,
chain_weight: torch.Tensor | None = None,
eps: float = 0.0,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
"""Fused embedding row gather (``embed_table[input_ids]``).
Requires the full vocab on-rank (replicated embedding). When
``chain_weight`` is given, also emits ``rmsnorm(gathered, chain_weight)``
(the first decoder layer's ``input_layernorm``) as a second output in the
same launch, so the returned pair is ``(residual, normed_input)``. Bit-exact
vs a plain gather followed by an ``RMSNorm``.
"""
assert embed_table.ndim == 2, embed_table.shape
ids = input_ids.view(-1)
(t,) = ids.shape
h = embed_table.shape[1]
if chain_weight is not None:
assert chain_weight.shape == (h,), (chain_weight.shape, h)
out = torch.empty((t, h), dtype=embed_table.dtype, device=embed_table.device)
normed = torch.empty_like(out) if chain_weight is not None else None
if t > 0:
block = triton.next_power_of_2(h)
_fused_embed_norm_kernel[(t,)](
ids,
embed_table,
embed_table.stride(0),
out,
normed if normed is not None else out,
chain_weight if chain_weight is not None else embed_table,
eps,
h,
block,
HAS_NORM=chain_weight is not None,
num_warps=min(32, max(4, block // 512)),
)
if normed is not None:
return out, normed
return out
@triton.jit
def _fused_embed_eh_norm_kernel(
pos_ptr,
ids_ptr, # [T] token ids
table_ptr, # [V, H] embedding table (full vocab, replicated on-rank)
table_stride,
prev_ptr, # [T, H] previous-step hidden
prev_stride,
enorm_w_ptr,
hnorm_w_ptr,
eps,
out_ptr, # [T, 2H]
out_stride,
H: tl.constexpr,
BLOCK: tl.constexpr,
):
"""MTP input fusion with a folded embedding gather: gather
``table[ids]``, zero it at position 0, RMSNorm(embed) with enorm and
RMSNorm(prev_hidden) with hnorm, written side-by-side into ``out`` ([N, 2H])
ready for the eh_proj GEMM. Replaces embedding lookup + where + 2x RMSNorm +
cat. Requires the full table on-rank (replicated embedding)."""
tok = tl.program_id(0)
off = tl.arange(0, BLOCK)
mask = off < H
pos = tl.load(pos_ptr + tok)
row = tl.load(ids_ptr + tok).to(tl.int64)
e = tl.load(table_ptr + row * table_stride + off, mask=mask, other=0.0)
e = tl.where(pos == 0, 0.0, e.to(tl.float32))
ew = tl.load(enorm_w_ptr + off, mask=mask)
e_normed = _rms_norm(e, ew, eps, H)
tl.store(out_ptr + tok * out_stride + off, e_normed, mask=mask)
p = tl.load(prev_ptr + tok * prev_stride + off, mask=mask, other=0.0)
hw = tl.load(hnorm_w_ptr + off, mask=mask)
p_normed = _rms_norm(p, hw, eps, H)
tl.store(out_ptr + tok * out_stride + H + off, p_normed, mask=mask)
# MTP fusion
def fused_embed_eh_norm(
positions: torch.Tensor,
input_ids: torch.Tensor,
embed_table: torch.Tensor,
previous_hidden: torch.Tensor,
enorm_w: torch.Tensor,
hnorm_w: torch.Tensor,
eps: float,
) -> torch.Tensor:
"""Fused ``cat([enorm(masked embed_table[ids]), hnorm(prev_hidden)])`` -> [N, 2H].
Folds the embedding row gather into the MTP eh-norm launch; requires the full
table on-rank (replicated embedding). Bit-exact vs gathering ``embed_table[
input_ids]`` and passing it to the model-local ``fused_eh_norm``.
"""
assert previous_hidden.ndim == 2 and embed_table.ndim == 2
n, h = previous_hidden.shape
assert positions.shape == (n,) and input_ids.view(-1).shape == (n,)
assert embed_table.shape[1] == h, (embed_table.shape, h)
assert enorm_w.shape == (h,) and hnorm_w.shape == (h,)
out = torch.empty(
n, 2 * h, dtype=previous_hidden.dtype, device=previous_hidden.device
)
_fused_embed_eh_norm_kernel[(n,)](
positions,
input_ids,
embed_table,
embed_table.stride(0),
previous_hidden,
previous_hidden.stride(0),
enorm_w,
hnorm_w,
eps,
out,
out.stride(0),
h,
triton.next_power_of_2(h),
)
return out
vllm/models/deepseek_v32/nvidia/model.py
+32 -6
View File
@@ -10,13 +10,15 @@ import vllm.envs as envs
from vllm.config import VllmConfig
from vllm.distributed import get_pp_group
from vllm.forward_context import get_forward_context, is_forward_context_available
from vllm.model_executor.layers.fused_embed_norm import (
fused_embed_norm,
has_full_vocab_on_rank,
make_input_embedding,
)
from vllm.model_executor.layers.fused_moe import (
fused_moe_make_expert_params_mapping,
)
from vllm.model_executor.layers.layernorm import RMSNorm
from vllm.model_executor.layers.vocab_parallel_embedding import (
VocabParallelEmbedding,
)
from vllm.model_executor.model_loader.weight_utils import (
default_weight_loader,
maybe_remap_kv_scale_name,
@@ -116,12 +118,18 @@ class DeepseekV32DecoderLayer(torch.nn.Module):
positions: torch.Tensor,
hidden_states: torch.Tensor,
residual: torch.Tensor | None,
attn_in: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
full_num_tokens = positions.shape[0]
if residual is None:
# First layer: hidden_states is the embedding (the residual).
residual = hidden_states
hidden_states = self.input_layernorm(hidden_states)
# ``attn_in`` is input_layernorm(embedding) already computed fused
# with the embedding gather; otherwise apply it here.
hidden_states = (
attn_in if attn_in is not None else self.input_layernorm(hidden_states)
)
elif self.use_sequence_parallel:
hidden_states, residual = self.input_layernorm(hidden_states, residual)
else:
@@ -182,14 +190,17 @@ class DeepseekV32Model(torch.nn.Module):
)
if get_pp_group().is_first_rank:
self.embed_tokens = VocabParallelEmbedding(
self.embed_tokens = make_input_embedding(
config.vocab_size,
config.hidden_size,
quant_config=quant_config,
prefix=f"{prefix}.embed_tokens",
tie_word_embeddings=getattr(config, "tie_word_embeddings", False),
)
else:
self.embed_tokens = PPMissingLayer()
# The fused embed+norm gather needs the full table on-rank.
self.replicated_embed = has_full_vocab_on_rank(self.embed_tokens)
self.start_layer, self.end_layer, self.layers = make_layers(
config.num_hidden_layers,
@@ -222,9 +233,21 @@ class DeepseekV32Model(torch.nn.Module):
intermediate_tensors: IntermediateTensors | None = None,
inputs_embeds: torch.Tensor | None = None,
) -> torch.Tensor | IntermediateTensors:
attn_in = None
if get_pp_group().is_first_rank:
if inputs_embeds is not None:
hidden_states = inputs_embeds
elif self.replicated_embed:
assert input_ids is not None
# Full table on-rank: gather the embedding and the first layer's
# input_layernorm in one launch. ``attn_in`` is the pre-normed
# attention input; ``hidden_states`` is the (residual) embedding.
hidden_states, attn_in = fused_embed_norm(
input_ids,
self.embed_tokens.weight,
chain_weight=self.layers[self.start_layer].input_layernorm.weight,
eps=self.config.rms_norm_eps,
)
else:
assert input_ids is not None
hidden_states = self.embed_input_ids(input_ids)
@@ -242,6 +265,8 @@ class DeepseekV32Model(torch.nn.Module):
forward_context.is_padding, hidden_states
)
hidden_states = sp_shard(hidden_states)
if attn_in is not None:
attn_in = sp_shard(attn_in)
assert residual is None, "Currently, SP is not supported with PP"
aux_hidden_states = []
@@ -253,7 +278,8 @@ class DeepseekV32Model(torch.nn.Module):
aux_hidden_states.append(
hidden_states if residual is None else hidden_states + residual
)
hidden_states, residual = layer(positions, hidden_states, residual)
hidden_states, residual = layer(positions, hidden_states, residual, attn_in)
attn_in = None
if not get_pp_group().is_last_rank:
assert not self.use_sequence_parallel, (
vllm/models/deepseek_v32/nvidia/mtp.py
+43 -15
View File
@@ -10,14 +10,16 @@ import vllm.envs as envs
from vllm._aiter_ops import rocm_aiter_ops
from vllm.config import VllmConfig
from vllm.forward_context import get_forward_context, is_forward_context_available
from vllm.model_executor.layers.fused_embed_norm import (
fused_embed_eh_norm,
has_full_vocab_on_rank,
make_input_embedding,
)
from vllm.model_executor.layers.fused_moe import (
fused_moe_make_expert_params_mapping,
)
from vllm.model_executor.layers.layernorm import RMSNorm
from vllm.model_executor.layers.logits_processor import LogitsProcessor
from vllm.model_executor.layers.vocab_parallel_embedding import (
VocabParallelEmbedding,
)
from vllm.model_executor.model_loader.mtp_validation import (
is_mtp_completeness_check_enabled,
)
@@ -93,18 +95,33 @@ class DeepseekV32MultiTokenPredictorLayer(nn.Module):
positions: torch.Tensor,
previous_hidden_states: torch.Tensor,
inputs_embeds: torch.Tensor | None = None,
embed_table: torch.Tensor | None = None,
spec_step_index: int = 0,
) -> torch.Tensor:
assert inputs_embeds is not None
# Fused: zero pos-0 embeds + enorm(embeds) + hnorm(prev) + cat -> [N, 2H].
eh_input = fused_eh_norm(
positions,
inputs_embeds,
previous_hidden_states,
self.enorm.weight,
self.hnorm.weight,
self.enorm.variance_epsilon,
)
# Fused zero pos-0 + enorm(embeds) + hnorm(prev) + cat -> [N, 2H]. With a
# replicated table the caller passes ``embed_table`` so the embedding
# lookup is folded in too (fused_embed_eh_norm); otherwise the embeds are
# precomputed and go through the model-local fused_eh_norm.
if embed_table is not None:
eh_input = fused_embed_eh_norm(
positions,
input_ids,
embed_table,
previous_hidden_states,
self.enorm.weight,
self.hnorm.weight,
self.enorm.variance_epsilon,
)
else:
assert inputs_embeds is not None
eh_input = fused_eh_norm(
positions,
inputs_embeds,
previous_hidden_states,
self.enorm.weight,
self.hnorm.weight,
self.enorm.variance_epsilon,
)
is_sequence_parallel = self.mtp_block.use_sequence_parallel
if is_sequence_parallel:
if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available():
@@ -158,11 +175,15 @@ class DeepseekV32MultiTokenPredictor(nn.Module):
)
}
)
self.embed_tokens = VocabParallelEmbedding(
self.embed_tokens = make_input_embedding(
config.vocab_size,
config.hidden_size,
quant_config=vllm_config.quant_config,
prefix=maybe_prefix(prefix, "embed_tokens"),
tie_word_embeddings=getattr(config, "tie_word_embeddings", False),
)
# A full on-rank table lets the eh_norm fusion fold in the embedding gather.
self.replicated_embed = has_full_vocab_on_rank(self.embed_tokens)
self.logits_processor = LogitsProcessor(config.vocab_size)
def set_skip_topk(self, skip: bool):
@@ -192,14 +213,21 @@ class DeepseekV32MultiTokenPredictor(nn.Module):
inputs_embeds: torch.Tensor | None = None,
spec_step_idx: int = 0,
) -> torch.Tensor:
# With a replicated table, defer the embedding gather to fused_eh_norm
# (folded into the enorm/hnorm/cat launch); otherwise gather it here.
embed_table = None
if inputs_embeds is None:
inputs_embeds = self.embed_tokens(input_ids)
if self.replicated_embed:
embed_table = self.embed_tokens.weight
else:
inputs_embeds = self.embed_tokens(input_ids)
current_step_idx = spec_step_idx % self.num_mtp_layers
return self.layers[str(self.mtp_start_layer_idx + current_step_idx)](
input_ids,
positions,
previous_hidden_states,
inputs_embeds,
embed_table,
current_step_idx,
)