diff --git a/tests/kernels/core/test_fused_embed_norm.py b/tests/kernels/core/test_fused_embed_norm.py new file mode 100644 index 0000000000..9da43dba83 --- /dev/null +++ b/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.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) diff --git a/vllm/envs.py b/vllm/envs.py index 15705b7c17..38b64f8c38 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -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 diff --git a/vllm/model_executor/layers/fused_embed_norm.py b/vllm/model_executor/layers/fused_embed_norm.py new file mode 100644 index 0000000000..ddeb5156ae --- /dev/null +++ b/vllm/model_executor/layers/fused_embed_norm.py @@ -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 diff --git a/vllm/models/deepseek_v32/nvidia/model.py b/vllm/models/deepseek_v32/nvidia/model.py index 0173d1ba6c..76aa6ca8e1 100644 --- a/vllm/models/deepseek_v32/nvidia/model.py +++ b/vllm/models/deepseek_v32/nvidia/model.py @@ -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, ( diff --git a/vllm/models/deepseek_v32/nvidia/mtp.py b/vllm/models/deepseek_v32/nvidia/mtp.py index fe2085af11..c2dcbdd64d 100644 --- a/vllm/models/deepseek_v32/nvidia/mtp.py +++ b/vllm/models/deepseek_v32/nvidia/mtp.py @@ -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, )