| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Replicated input embedding + its fused gather/norm kernels. |
| 4 | |
| 5 | Groups the ``VLLM_REPLICATE_EMBED`` path in one place: the embedding factory, |
| 6 | the predicate that says whether the fusions apply, and the two Triton fusions |
| 7 | the full on-rank table unlocks -- |
| 8 | |
| 9 | * ``fused_embed_norm``: gather + a chained RMSNorm (e.g. the first decoder |
| 10 | layer's ``input_layernorm``), and |
| 11 | * ``fused_embed_eh_norm``: gather + pos-0 zeroing + enorm/hnorm + cat, the |
| 12 | embed/previous-hidden input norm for a speculative (MTP/eagle) depth layer |
| 13 | (the replicated-table analogue of the model-local ``fused_eh_norm``, which |
| 14 | takes precomputed embeds). |
| 15 | |
| 16 | Self-contained (no model-local imports) so it can live under ``layers/``. |
| 17 | """ |
| 18 | |
| 19 | import torch |
| 20 | |
| 21 | import vllm.envs as envs |
| 22 | from vllm.distributed import get_tensor_model_parallel_world_size |
| 23 | from vllm.model_executor.layers.vocab_parallel_embedding import ( |
| 24 | UnquantizedEmbeddingMethod, |
| 25 | VocabParallelEmbedding, |
| 26 | ) |
| 27 | from vllm.triton_utils import tl, triton |
| 28 | |
| 29 | |
| 30 | @triton.jit |
| 31 | def _rms_norm(x, w, eps, HIDDEN_SIZE: tl.constexpr): |
| 32 | x = x.to(tl.float32) |
| 33 | mean_sq = tl.sum(x * x, axis=0) / HIDDEN_SIZE |
| 34 | rrms = tl.rsqrt(mean_sq + eps) |
| 35 | w = w.to(tl.float32) |
| 36 | return (x * rrms) * w |
| 37 | |
| 38 | |
| 39 | def make_input_embedding( |
| 40 | num_embeddings: int, |
| 41 | embedding_dim: int, |
| 42 | *, |
| 43 | params_dtype: torch.dtype | None = None, |
| 44 | quant_config=None, |
| 45 | prefix: str = "", |
| 46 | tie_word_embeddings: bool = False, |
| 47 | ) -> VocabParallelEmbedding: |
| 48 | """Input token embedding with an optional replicated escape hatch. |
| 49 | |
| 50 | ``VLLM_REPLICATE_EMBED=1`` builds the embedding with ``disable_tp``: the full |
| 51 | table lives on every rank and the lookup is a local gather with no mask and |
| 52 | no all-reduce, which unlocks the fused gather+norm path. The cost is a full |
| 53 | table per rank at TP>1 (no extra memory at TP=1, where vocab-parallel is |
| 54 | already unsharded). A replicated, unsharded table cannot be tied to a |
| 55 | vocab-parallel ``ParallelLMHead``, so tied word embeddings are rejected at |
| 56 | TP>1 (at TP=1 ``disable_tp`` is a no-op and tying still works). |
| 57 | """ |
| 58 | disable_tp = envs.VLLM_REPLICATE_EMBED |
| 59 | if disable_tp and tie_word_embeddings: |
| 60 | assert get_tensor_model_parallel_world_size() == 1, ( |
| 61 | "VLLM_REPLICATE_EMBED is unsupported with tied word embeddings " |
| 62 | "(the replicated table cannot tie to a vocab-parallel lm_head)" |
| 63 | ) |
| 64 | return VocabParallelEmbedding( |
| 65 | num_embeddings, |
| 66 | embedding_dim, |
| 67 | params_dtype=params_dtype, |
| 68 | quant_config=quant_config, |
| 69 | prefix=prefix, |
| 70 | disable_tp=disable_tp, |
| 71 | ) |
| 72 | |
| 73 | |
| 74 | def has_full_vocab_on_rank(embedding: torch.nn.Module) -> bool: |
| 75 | """Whether ``embedding.weight`` is the whole vocab as a plain [V, H] table. |
| 76 | |
| 77 | The fused gather kernels index the table directly, so they need every row |
| 78 | on-rank (``disable_tp``, or any TP=1 run) and an unquantized weight. |
| 79 | """ |
| 80 | return getattr(embedding, "tp_size", 0) == 1 and isinstance( |
| 81 | getattr(embedding, "quant_method", None), UnquantizedEmbeddingMethod |
| 82 | ) |
| 83 | |
| 84 | |
| 85 | @triton.jit |
| 86 | def _fused_embed_norm_kernel( |
| 87 | ids_ptr, # [T] token ids |
| 88 | table_ptr, # [V, H] embedding table (full vocab, replicated on-rank) |
| 89 | table_stride_0, |
| 90 | out_ptr, # [T, H] gathered embedding (the residual stream) |
| 91 | normed_ptr, # [T, H] rmsnorm(out, chain_w) (HAS_NORM only) |
| 92 | chain_w_ptr, # [H] next norm weight (HAS_NORM only) |
| 93 | eps, |
| 94 | H: tl.constexpr, |
| 95 | BLOCK: tl.constexpr, |
| 96 | HAS_NORM: tl.constexpr, |
| 97 | ): |
| 98 | tok = tl.program_id(0).to(tl.int64) |
| 99 | off = tl.arange(0, BLOCK) |
| 100 | mask = off < H |
| 101 | row = tl.load(ids_ptr + tok).to(tl.int64) |
| 102 | x = tl.load(table_ptr + row * table_stride_0 + off, mask=mask, other=0.0) |
| 103 | tl.store(out_ptr + tok * H + off, x, mask=mask) |
| 104 | if HAS_NORM: |
| 105 | w = tl.load(chain_w_ptr + off, mask=mask) |
| 106 | y = _rms_norm(x, w, eps, H).to(normed_ptr.dtype.element_ty) |
| 107 | tl.store(normed_ptr + tok * H + off, y, mask=mask) |
| 108 | |
| 109 | |
| 110 | # Base model fusion |
| 111 | def fused_embed_norm( |
| 112 | input_ids: torch.Tensor, |
| 113 | embed_table: torch.Tensor, |
| 114 | chain_weight: torch.Tensor | None = None, |
| 115 | eps: float = 0.0, |
| 116 | ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: |
| 117 | """Fused embedding row gather (``embed_table[input_ids]``). |
| 118 | |
| 119 | Requires the full vocab on-rank (replicated embedding). When |
| 120 | ``chain_weight`` is given, also emits ``rmsnorm(gathered, chain_weight)`` |
| 121 | (the first decoder layer's ``input_layernorm``) as a second output in the |
| 122 | same launch, so the returned pair is ``(residual, normed_input)``. Bit-exact |
| 123 | vs a plain gather followed by an ``RMSNorm``. |
| 124 | """ |
| 125 | assert embed_table.ndim == 2, embed_table.shape |
| 126 | ids = input_ids.view(-1) |
| 127 | (t,) = ids.shape |
| 128 | h = embed_table.shape[1] |
| 129 | if chain_weight is not None: |
| 130 | assert chain_weight.shape == (h,), (chain_weight.shape, h) |
| 131 | out = torch.empty((t, h), dtype=embed_table.dtype, device=embed_table.device) |
| 132 | normed = torch.empty_like(out) if chain_weight is not None else None |
| 133 | if t > 0: |
| 134 | block = triton.next_power_of_2(h) |
| 135 | _fused_embed_norm_kernel[(t,)]( |
| 136 | ids, |
| 137 | embed_table, |
| 138 | embed_table.stride(0), |
| 139 | out, |
| 140 | normed if normed is not None else out, |
| 141 | chain_weight if chain_weight is not None else embed_table, |
| 142 | eps, |
| 143 | h, |
| 144 | block, |
| 145 | HAS_NORM=chain_weight is not None, |
| 146 | num_warps=min(32, max(4, block // 512)), |
| 147 | ) |
| 148 | if normed is not None: |
| 149 | return out, normed |
| 150 | return out |
| 151 | |
| 152 | |
| 153 | @triton.jit |
| 154 | def _fused_embed_eh_norm_kernel( |
| 155 | pos_ptr, |
| 156 | ids_ptr, # [T] token ids |
| 157 | table_ptr, # [V, H] embedding table (full vocab, replicated on-rank) |
| 158 | table_stride, |
| 159 | prev_ptr, # [T, H] previous-step hidden |
| 160 | prev_stride, |
| 161 | enorm_w_ptr, |
| 162 | hnorm_w_ptr, |
| 163 | eps, |
| 164 | out_ptr, # [T, 2H] |
| 165 | out_stride, |
| 166 | H: tl.constexpr, |
| 167 | BLOCK: tl.constexpr, |
| 168 | ): |
| 169 | """MTP input fusion with a folded embedding gather: gather |
| 170 | ``table[ids]``, zero it at position 0, RMSNorm(embed) with enorm and |
| 171 | RMSNorm(prev_hidden) with hnorm, written side-by-side into ``out`` ([N, 2H]) |
| 172 | ready for the eh_proj GEMM. Replaces embedding lookup + where + 2x RMSNorm + |
| 173 | cat. Requires the full table on-rank (replicated embedding).""" |
| 174 | tok = tl.program_id(0) |
| 175 | off = tl.arange(0, BLOCK) |
| 176 | mask = off < H |
| 177 | |
| 178 | pos = tl.load(pos_ptr + tok) |
| 179 | row = tl.load(ids_ptr + tok).to(tl.int64) |
| 180 | e = tl.load(table_ptr + row * table_stride + off, mask=mask, other=0.0) |
| 181 | e = tl.where(pos == 0, 0.0, e.to(tl.float32)) |
| 182 | ew = tl.load(enorm_w_ptr + off, mask=mask) |
| 183 | e_normed = _rms_norm(e, ew, eps, H) |
| 184 | tl.store(out_ptr + tok * out_stride + off, e_normed, mask=mask) |
| 185 | |
| 186 | p = tl.load(prev_ptr + tok * prev_stride + off, mask=mask, other=0.0) |
| 187 | hw = tl.load(hnorm_w_ptr + off, mask=mask) |
| 188 | p_normed = _rms_norm(p, hw, eps, H) |
| 189 | tl.store(out_ptr + tok * out_stride + H + off, p_normed, mask=mask) |
| 190 | |
| 191 | |
| 192 | # MTP fusion |
| 193 | def fused_embed_eh_norm( |
| 194 | positions: torch.Tensor, |
| 195 | input_ids: torch.Tensor, |
| 196 | embed_table: torch.Tensor, |
| 197 | previous_hidden: torch.Tensor, |
| 198 | enorm_w: torch.Tensor, |
| 199 | hnorm_w: torch.Tensor, |
| 200 | eps: float, |
| 201 | ) -> torch.Tensor: |
| 202 | """Fused ``cat([enorm(masked embed_table[ids]), hnorm(prev_hidden)])`` -> [N, 2H]. |
| 203 | |
| 204 | Folds the embedding row gather into the MTP eh-norm launch; requires the full |
| 205 | table on-rank (replicated embedding). Bit-exact vs gathering ``embed_table[ |
| 206 | input_ids]`` and passing it to the model-local ``fused_eh_norm``. |
| 207 | """ |
| 208 | assert previous_hidden.ndim == 2 and embed_table.ndim == 2 |
| 209 | n, h = previous_hidden.shape |
| 210 | assert positions.shape == (n,) and input_ids.view(-1).shape == (n,) |
| 211 | assert embed_table.shape[1] == h, (embed_table.shape, h) |
| 212 | assert enorm_w.shape == (h,) and hnorm_w.shape == (h,) |
| 213 | out = torch.empty( |
| 214 | n, 2 * h, dtype=previous_hidden.dtype, device=previous_hidden.device |
| 215 | ) |
| 216 | _fused_embed_eh_norm_kernel[(n,)]( |
| 217 | positions, |
| 218 | input_ids, |
| 219 | embed_table, |
| 220 | embed_table.stride(0), |
| 221 | previous_hidden, |
| 222 | previous_hidden.stride(0), |
| 223 | enorm_w, |
| 224 | hnorm_w, |
| 225 | eps, |
| 226 | out, |
| 227 | out.stride(0), |
| 228 | h, |
| 229 | triton.next_power_of_2(h), |
| 230 | ) |
| 231 | return out |