mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-26 16:35:05 +00:00
31 lines
883 B
Python
31 lines
883 B
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# Derived from vLLM; modified by DeepSeek.
|
|
|
|
import torch
|
|
|
|
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:
|
|
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.")
|