Replicated embedding and norm fusion for DSV3 flat model (#48484)
Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>
Jee Jee Li committed
Aug 18, 2026 at 14:53 UTC
2687fec6ef744f7d8ac174303e9a2ee6fe8d7153
5 files changed
+385
-21
tests/kernels/core/test_fused_embed_norm.py
new
+75
@@ -0,0 +1,75 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Unit tests for the replicated-embedding fused gather/norm kernels
4
+(``vllm.model_executor.layers.fused_embed_norm``).
5
+
6
+The guarantee: enabling ``VLLM_REPLICATE_EMBED`` (replicated table + fused
7
+kernels) must not change model outputs. The gathered residual is bit-exact and
8
+the fused norms match the unfused reference.
9
+"""
10
+
11
+import pytest
12
+import torch
13
+
14
+from vllm.model_executor.layers.fused_embed_norm import (
15
+ fused_embed_eh_norm,
16
+ fused_embed_norm,
17
+)
18
+
19
+# The model-local (untouched) eh-norm the replicate path must match.
20
+from vllm.models.deepseek_v32.common.kernels import fused_eh_norm
21
+from vllm.platforms import current_platform
22
+from vllm.utils.torch_utils import set_random_seed
23
+
24
+DTYPE = torch.bfloat16
25
+VOCAB, HIDDEN, NUM_TOKENS, EPS = 8192, 4096, 129, 1e-6
26
+
27
+requires_cuda = pytest.mark.skipif(
28
+ not current_platform.is_cuda_alike(),
29
+ reason="fused embed/norm Triton kernels require a CUDA/ROCm device",
30
+)
31
+
32
+
33
+def _rmsnorm(x: torch.Tensor, w: torch.Tensor, eps: float) -> torch.Tensor:
34
+ # Full-precision (fp32) reference RMSNorm.
35
+ var = x.float().pow(2).mean(dim=-1, keepdim=True)
36
+ return x.float() * torch.rsqrt(var + eps) * w.float()
37
+
38
+
39
+@requires_cuda
40
+@torch.inference_mode()
41
+def test_fused_embed_norm_matches_reference():
42
+ """Main-model fusion: the residual is the exact gather and the second output
43
+ is a correct RMSNorm. The norm matches a full-precision reference to ~2 bf16
44
+ ulp (rtol 1e-2) -- that gap is bf16 rounding, not the kernel."""
45
+ set_random_seed(13)
46
+ table = torch.randn(VOCAB, HIDDEN, dtype=DTYPE, device="cuda")
47
+ ids = torch.randint(0, VOCAB, (NUM_TOKENS,), dtype=torch.int32, device="cuda")
48
+ weight = torch.empty(HIDDEN, dtype=DTYPE, device="cuda").normal_(1.0, 0.1)
49
+
50
+ residual, normed = fused_embed_norm(ids, table, chain_weight=weight, eps=EPS)
51
+
52
+ embeds = table[ids.long()]
53
+ torch.testing.assert_close(residual, embeds, atol=0.0, rtol=0.0)
54
+ torch.testing.assert_close(
55
+ normed.float(), _rmsnorm(embeds, weight, EPS), atol=1e-3, rtol=1e-2
56
+ )
57
+
58
+
59
+@requires_cuda
60
+@torch.inference_mode()
61
+def test_fused_embed_eh_norm_matches_reference():
62
+ """MTP fusion (folded gather) is bit-exact vs gathering the embeds and
63
+ feeding the untouched model-local ``fused_eh_norm``."""
64
+ set_random_seed(13)
65
+ table = torch.randn(VOCAB, HIDDEN, dtype=DTYPE, device="cuda")
66
+ ids = torch.randint(0, VOCAB, (NUM_TOKENS,), dtype=torch.int32, device="cuda")
67
+ prev = torch.randn(NUM_TOKENS, HIDDEN, dtype=DTYPE, device="cuda")
68
+ enorm_w = torch.randn(HIDDEN, dtype=DTYPE, device="cuda")
69
+ hnorm_w = torch.randn(HIDDEN, dtype=DTYPE, device="cuda")
70
+ positions = torch.arange(NUM_TOKENS, device="cuda") # includes pos 0
71
+
72
+ fused = fused_embed_eh_norm(positions, ids, table, prev, enorm_w, hnorm_w, EPS)
73
+ ref = fused_eh_norm(positions, table[ids.long()], prev, enorm_w, hnorm_w, EPS)
74
+
75
+ torch.testing.assert_close(fused, ref, atol=0.0, rtol=0.0)
vllm/envs.py
+4
@@ -156,6 +156,7 @@ if TYPE_CHECKING:
156
VLLM_ENABLE_V1_MULTIPROCESSING: bool = True
157
VLLM_LOG_BATCHSIZE_INTERVAL: float = -1
158
VLLM_DISABLE_COMPILE_CACHE: bool = False
159
+ VLLM_REPLICATE_EMBED: bool = False
160
VLLM_USE_LAYERNAME: bool = True
161
Q_SCALE_CONSTANT: int = 200
162
K_SCALE_CONSTANT: int = 200
@@ -620,6 +621,9 @@ environment_variables: dict[str, Callable[[], Any]] = {
621
# Enable batch-invariant mode: deterministic results regardless of
622
# batch composition. Requires NVIDIA GPU with compute capability >= 9.0.
623
"VLLM_BATCH_INVARIANT": lambda: bool(int(os.getenv("VLLM_BATCH_INVARIANT", "0"))),
624
+ "VLLM_REPLICATE_EMBED": lambda: (
625
+ os.getenv("VLLM_REPLICATE_EMBED", "0").strip().lower() in ("1", "true")
626
+ ),
627
# Use tensor descriptors for Q/K/V loads and output stores in the
628
# Triton unified-attention kernel. Enables HW 2D block reads on
629
# Intel XPU; the non-TD branch is dead-code-eliminated at Triton
vllm/model_executor/layers/fused_embed_norm.py
new
+231
@@ -0,0 +1,231 @@
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
vllm/models/deepseek_v32/nvidia/model.py
+32
-6
@@ -10,13 +10,15 @@ import vllm.envs as envs
10
from vllm.config import VllmConfig
11
from vllm.distributed import get_pp_group
12
from vllm.forward_context import get_forward_context, is_forward_context_available
13
+from vllm.model_executor.layers.fused_embed_norm import (
14
+ fused_embed_norm,
15
+ has_full_vocab_on_rank,
16
+ make_input_embedding,
17
+)
18
from vllm.model_executor.layers.fused_moe import (
19
fused_moe_make_expert_params_mapping,
20
)
21
from vllm.model_executor.layers.layernorm import RMSNorm
17
-from vllm.model_executor.layers.vocab_parallel_embedding import (
18
- VocabParallelEmbedding,
19
-)
22
from vllm.model_executor.model_loader.weight_utils import (
23
default_weight_loader,
24
maybe_remap_kv_scale_name,
@@ -116,12 +118,18 @@ class DeepseekV32DecoderLayer(torch.nn.Module):
118
positions: torch.Tensor,
119
hidden_states: torch.Tensor,
120
residual: torch.Tensor | None,
121
+ attn_in: torch.Tensor | None = None,
122
) -> tuple[torch.Tensor, torch.Tensor]:
123
full_num_tokens = positions.shape[0]
124
125
if residual is None:
126
+ # First layer: hidden_states is the embedding (the residual).
127
residual = hidden_states
124
- hidden_states = self.input_layernorm(hidden_states)
128
+ # ``attn_in`` is input_layernorm(embedding) already computed fused
129
+ # with the embedding gather; otherwise apply it here.
130
+ hidden_states = (
131
+ attn_in if attn_in is not None else self.input_layernorm(hidden_states)
132
+ )
133
elif self.use_sequence_parallel:
134
hidden_states, residual = self.input_layernorm(hidden_states, residual)
135
else:
@@ -182,14 +190,17 @@ class DeepseekV32Model(torch.nn.Module):
190
)
191
192
if get_pp_group().is_first_rank:
185
- self.embed_tokens = VocabParallelEmbedding(
193
+ self.embed_tokens = make_input_embedding(
194
config.vocab_size,
195
config.hidden_size,
196
quant_config=quant_config,
197
prefix=f"{prefix}.embed_tokens",
198
+ tie_word_embeddings=getattr(config, "tie_word_embeddings", False),
199
)
200
else:
201
self.embed_tokens = PPMissingLayer()
202
+ # The fused embed+norm gather needs the full table on-rank.
203
+ self.replicated_embed = has_full_vocab_on_rank(self.embed_tokens)
204
205
self.start_layer, self.end_layer, self.layers = make_layers(
206
config.num_hidden_layers,
@@ -222,9 +233,21 @@ class DeepseekV32Model(torch.nn.Module):
233
intermediate_tensors: IntermediateTensors | None = None,
234
inputs_embeds: torch.Tensor | None = None,
235
) -> torch.Tensor | IntermediateTensors:
236
+ attn_in = None
237
if get_pp_group().is_first_rank:
238
if inputs_embeds is not None:
239
hidden_states = inputs_embeds
240
+ elif self.replicated_embed:
241
+ assert input_ids is not None
242
+ # Full table on-rank: gather the embedding and the first layer's
243
+ # input_layernorm in one launch. ``attn_in`` is the pre-normed
244
+ # attention input; ``hidden_states`` is the (residual) embedding.
245
+ hidden_states, attn_in = fused_embed_norm(
246
+ input_ids,
247
+ self.embed_tokens.weight,
248
+ chain_weight=self.layers[self.start_layer].input_layernorm.weight,
249
+ eps=self.config.rms_norm_eps,
250
+ )
251
else:
252
assert input_ids is not None
253
hidden_states = self.embed_input_ids(input_ids)
@@ -242,6 +265,8 @@ class DeepseekV32Model(torch.nn.Module):
265
forward_context.is_padding, hidden_states
266
)
267
hidden_states = sp_shard(hidden_states)
268
+ if attn_in is not None:
269
+ attn_in = sp_shard(attn_in)
270
assert residual is None, "Currently, SP is not supported with PP"
271
272
aux_hidden_states = []
@@ -253,7 +278,8 @@ class DeepseekV32Model(torch.nn.Module):
278
aux_hidden_states.append(
279
hidden_states if residual is None else hidden_states + residual
280
)
256
- hidden_states, residual = layer(positions, hidden_states, residual)
281
+ hidden_states, residual = layer(positions, hidden_states, residual, attn_in)
282
+ attn_in = None
283
284
if not get_pp_group().is_last_rank:
285
assert not self.use_sequence_parallel, (
vllm/models/deepseek_v32/nvidia/mtp.py
+43
-15
@@ -10,14 +10,16 @@ import vllm.envs as envs
10
from vllm._aiter_ops import rocm_aiter_ops
11
from vllm.config import VllmConfig
12
from vllm.forward_context import get_forward_context, is_forward_context_available
13
+from vllm.model_executor.layers.fused_embed_norm import (
14
+ fused_embed_eh_norm,
15
+ has_full_vocab_on_rank,
16
+ make_input_embedding,
17
+)
18
from vllm.model_executor.layers.fused_moe import (
19
fused_moe_make_expert_params_mapping,
20
)
21
from vllm.model_executor.layers.layernorm import RMSNorm
22
from vllm.model_executor.layers.logits_processor import LogitsProcessor
18
-from vllm.model_executor.layers.vocab_parallel_embedding import (
19
- VocabParallelEmbedding,
20
-)
23
from vllm.model_executor.model_loader.mtp_validation import (
24
is_mtp_completeness_check_enabled,
25
)
@@ -93,18 +95,33 @@ class DeepseekV32MultiTokenPredictorLayer(nn.Module):
95
positions: torch.Tensor,
96
previous_hidden_states: torch.Tensor,
97
inputs_embeds: torch.Tensor | None = None,
98
+ embed_table: torch.Tensor | None = None,
99
spec_step_index: int = 0,
100
) -> torch.Tensor:
98
- assert inputs_embeds is not None
99
- # Fused: zero pos-0 embeds + enorm(embeds) + hnorm(prev) + cat -> [N, 2H].
100
- eh_input = fused_eh_norm(
101
- positions,
102
- inputs_embeds,
103
- previous_hidden_states,
104
- self.enorm.weight,
105
- self.hnorm.weight,
106
- self.enorm.variance_epsilon,
107
- )
101
+ # Fused zero pos-0 + enorm(embeds) + hnorm(prev) + cat -> [N, 2H]. With a
102
+ # replicated table the caller passes ``embed_table`` so the embedding
103
+ # lookup is folded in too (fused_embed_eh_norm); otherwise the embeds are
104
+ # precomputed and go through the model-local fused_eh_norm.
105
+ if embed_table is not None:
106
+ eh_input = fused_embed_eh_norm(
107
+ positions,
108
+ input_ids,
109
+ embed_table,
110
+ previous_hidden_states,
111
+ self.enorm.weight,
112
+ self.hnorm.weight,
113
+ self.enorm.variance_epsilon,
114
+ )
115
+ else:
116
+ assert inputs_embeds is not None
117
+ eh_input = fused_eh_norm(
118
+ positions,
119
+ inputs_embeds,
120
+ previous_hidden_states,
121
+ self.enorm.weight,
122
+ self.hnorm.weight,
123
+ self.enorm.variance_epsilon,
124
+ )
125
is_sequence_parallel = self.mtp_block.use_sequence_parallel
126
if is_sequence_parallel:
127
if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available():
@@ -158,11 +175,15 @@ class DeepseekV32MultiTokenPredictor(nn.Module):
175
)
176
}
177
)
161
- self.embed_tokens = VocabParallelEmbedding(
178
+ self.embed_tokens = make_input_embedding(
179
config.vocab_size,
180
config.hidden_size,
181
+ quant_config=vllm_config.quant_config,
182
prefix=maybe_prefix(prefix, "embed_tokens"),
183
+ tie_word_embeddings=getattr(config, "tie_word_embeddings", False),
184
)
185
+ # A full on-rank table lets the eh_norm fusion fold in the embedding gather.
186
+ self.replicated_embed = has_full_vocab_on_rank(self.embed_tokens)
187
self.logits_processor = LogitsProcessor(config.vocab_size)
188
189
def set_skip_topk(self, skip: bool):
@@ -192,14 +213,21 @@ class DeepseekV32MultiTokenPredictor(nn.Module):
213
inputs_embeds: torch.Tensor | None = None,
214
spec_step_idx: int = 0,
215
) -> torch.Tensor:
216
+ # With a replicated table, defer the embedding gather to fused_eh_norm
217
+ # (folded into the enorm/hnorm/cat launch); otherwise gather it here.
218
+ embed_table = None
219
if inputs_embeds is None:
196
- inputs_embeds = self.embed_tokens(input_ids)
220
+ if self.replicated_embed:
221
+ embed_table = self.embed_tokens.weight
222
+ else:
223
+ inputs_embeds = self.embed_tokens(input_ids)
224
current_step_idx = spec_step_idx % self.num_mtp_layers
225
return self.layers[str(self.mtp_start_layer_idx + current_step_idx)](
226
input_ids,
227
positions,
228
previous_hidden_states,
229
inputs_embeds,
230
+ embed_table,
231
current_step_idx,
232
)
233