Init

Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>

tests/kernels/core/test_fused_embed_norm.py

@@ -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.nvidia.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

@@ -147,6 +147,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
@@ -583,6 +584,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 Xe2/Xe3; the non-TD branch is dead-code-eliminated at Triton

vllm/model_executor/layers/fused_embed_norm.py

@@ -0,0 +1,244 @@
+# 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 replicated embedding
+module, its factory (with the fallback to vocab-parallel), 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.model_executor.layers.vocab_parallel_embedding import (
+    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
+
+
+class ReplicatedEmbedding(torch.nn.Embedding):
+    """Fully-replicated input token embedding for GLM-5.2 / DeepSeek-V32.
+
+    The full [num_embeddings, embedding_dim] table lives on every TP rank and
+    the forward is a local lookup with NO all-reduce (unlike
+    VocabParallelEmbedding, which shards the vocab and all-reduces the output).
+    The full on-rank table also enables the fused gather+norm kernels
+    (``fused_embed_norm`` / ``fused_eh_norm`` GATHER mode). The weight has no
+    ``weight_loader`` attr, so it loads via ``default_weight_loader`` (a
+    shape-checked full-tensor copy). The only addition over
+    ``torch.nn.Embedding`` is the int32->int64 index cast that ``F.embedding``
+    requires (vLLM feeds int32 input_ids).
+    """
+
+    def forward(self, input_: torch.Tensor) -> torch.Tensor:
+        return super().forward(input_.long())
+
+
+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,
+):
+    """Input token embedding with an optional replicated escape hatch.
+
+    With ``VLLM_REPLICATE_EMBED=1`` use a fully-replicated ``ReplicatedEmbedding``
+    to unlock the fused gather+norm path (and, at TP>1, skip the embedding
+    all-reduce), at the cost of a full table per rank at TP>1 (no extra memory at
+    TP=1, where vocab-parallel is already unsharded). The replicated table is
+    always the raw (unquantized) ``params_dtype``; ``quant_config`` is accepted
+    only for the vocab-parallel fallback (embeddings are unquantized regardless).
+    Assumes an untied embedding -- a replicated, unsharded table cannot be tied
+    to a vocab-parallel ``ParallelLMHead``, so tied word embeddings are rejected.
+    Otherwise falls back to the default TP-sharded ``VocabParallelEmbedding``
+    (byte-identical to before).
+    """
+    if envs.VLLM_REPLICATE_EMBED:
+        assert not tie_word_embeddings, (
+            "VLLM_REPLICATE_EMBED is unsupported with tied word embeddings "
+            "(the replicated table cannot tie to a vocab-parallel lm_head)"
+        )
+        emb = ReplicatedEmbedding(
+            num_embeddings,
+            embedding_dim,
+            dtype=params_dtype or torch.get_default_dtype(),
+        )
+        emb.weight.requires_grad_(False)
+        return emb
+    return VocabParallelEmbedding(
+        num_embeddings,
+        embedding_dim,
+        params_dtype=params_dtype,
+        quant_config=quant_config,
+        prefix=prefix,
+    )
+
+
+@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

@@ -8,13 +8,15 @@ import torch
 
 from vllm.config import VllmConfig
 from vllm.distributed import get_pp_group
+from vllm.model_executor.layers.fused_embed_norm import (
+    ReplicatedEmbedding,
+    fused_embed_norm,
+    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,
@@ -103,11 +105,16 @@ 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]:
         if residual is None:
-            # First layer: hidden_states is the (already reduced) embedding.
+            # 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)
+            )
         else:
             # The previous layer's MLP/MoE output is left un-reduced; fuse its
             # all-reduce into this input_layernorm.
@@ -151,14 +158,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 = isinstance(self.embed_tokens, ReplicatedEmbedding)
 
         self.start_layer, self.end_layer, self.layers = make_layers(
             config.num_hidden_layers,
@@ -193,9 +203,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)
@@ -212,7 +234,8 @@ class DeepseekV32Model(torch.nn.Module):
         ):
             if idx in self.aux_hidden_state_layers:
                 aux_hidden_states.append(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:
             return IntermediateTensors(

vllm/models/deepseek_v32/nvidia/mtp.py

@@ -9,14 +9,16 @@ import torch.nn as nn
 from vllm._aiter_ops import rocm_aiter_ops
 from vllm.config import VllmConfig
 from vllm.distributed import tensor_model_parallel_all_reduce
+from vllm.model_executor.layers.fused_embed_norm import (
+    ReplicatedEmbedding,
+    fused_embed_eh_norm,
+    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.weight_utils import (
     default_weight_loader,
     maybe_remap_kv_scale_name,
@@ -73,18 +75,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,
+            )
         hidden_states = self.eh_proj(eh_input)
         hidden_states, residual = self.mtp_block(
             positions=positions, hidden_states=hidden_states, residual=None
@@ -124,11 +141,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 replicated table lets the eh_norm fusion fold in the embedding gather.
+        self.replicated_embed = isinstance(self.embed_tokens, ReplicatedEmbedding)
         self.logits_processor = LogitsProcessor(config.vocab_size)
 
     def set_skip_topk(self, skip: bool):
@@ -158,14 +179,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,
         )