mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-29 23:37:28 +08:00
0016315f80
Both drain the device stream in torch-neuronx (a cast round-trips through the host, a scalar operand is a host->device copy): int32 ids index directly and out-of-range rows are cleared with masked_fill. Tiny cp=2 nextn=5 on device: 9/9 golden tokens. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
45 lines
1.9 KiB
Python
45 lines
1.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# Derived from vLLM; modified by DeepSeek.
|
|
|
|
import torch
|
|
|
|
from deepinfer import PLATFORM
|
|
from deepinfer.model_executor.layers.tilelang_ops.gather import gather
|
|
from deepinfer.model_executor.utils import reconstruct_param
|
|
|
|
|
|
class VocabParallelEmbedding(torch.nn.Module):
|
|
def __init__(
|
|
self,
|
|
num_embeddings: int,
|
|
embedding_dim: int,
|
|
weight: torch.nn.Parameter,
|
|
) -> None:
|
|
super().__init__()
|
|
|
|
assert weight.shape == (num_embeddings, embedding_dim)
|
|
self.weight = reconstruct_param(weight)
|
|
|
|
def forward(self, input_: torch.Tensor) -> torch.Tensor:
|
|
if self.weight.device != input_.device:
|
|
# Host-resident table (Neuron, see DeepseekForCausalLM.load_model): move the looked-up rows only.
|
|
return gather(self.weight, input_.cpu()).to(input_.device)
|
|
if PLATFORM.is_neuron and not torch.compiler.is_compiling():
|
|
# Eagerly the zero-filling gather is ~5 NEFF launches, and its scalar operand is a host->device copy,
|
|
# which drains the stream. One lookup instead; out-of-range ids still read zeros unless the caller's
|
|
# ids are always valid (ids_in_range: DSpark's Markov embedding of sampled draft tokens).
|
|
# No dtype casts either: eagerly they round-trip through the host (int32 ids index directly).
|
|
n = self.weight.shape[0]
|
|
rows = torch.nn.functional.embedding(input_.clamp(0, n - 1), self.weight)
|
|
if getattr(self, "ids_in_range", False):
|
|
return rows
|
|
return rows.masked_fill(((input_ < 0) | (input_ >= n)).unsqueeze(-1), 0)
|
|
output = gather(self.weight, input_)
|
|
return output
|
|
|
|
|
|
class ParallelLMHead(VocabParallelEmbedding):
|
|
def forward(self, input_: torch.Tensor) -> None:
|
|
del input_
|
|
raise RuntimeError("LMHead's weights should be used in the sampler.")
|