mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-31 00:36:20 +08:00
c24574bfd6
- Layer-major KV / state layout so per-layer views are contiguous graph inputs (the executor copied non-contiguous inputs every call). - One compiled entry per (class, method), rank as a registered buffer, no layer_idx guard: every layer reuses the same graphs and loaded NEFFs. - fp8 weight rescale on the host (numpy LUT, ~2 min instead of 35+ min of eager per-shape compiles); engram sf packing on the host. - int16 ue8m0 packing in fp32 (int32 interleaved slices narrowed to int16 crash neuronx-cc, NCC_IBCG902). - Token embedding and the text model's vision tower stay in host memory on Neuron (2.5 GiB of HBM a rank). - offline example: --max_num_seqs. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
34 lines
1.1 KiB
Python
34 lines
1.1 KiB
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:
|
|
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)
|
|
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.")
|