f355fdd39e
CP shared KV with HiCache and EAGLE needs host backup to overlap forward while keeping radix visibility synchronous. The change reserves host slots before forward, drives target and draft backup from explicit layer-end hooks, and commits host visibility only after the final target/draft ack. It also probes the final insertion prefix before early reservation so repeated EAGLE prompts do not prepare duplicate suffix backups that later rollback as insert_miss. Constraint: CP ranks use independent shared-KV pools, so target/draft host state must remain atomically visible at the radix boundary. Constraint: Fused MLA and NSA store paths can bypass store-side notifier hooks, so layer end is the safer backup progress boundary. Rejected: Store-side backup notifier as the primary trigger | fused store and zero-local paths made notifier coverage fragile. Rejected: Reserve from cache_protected_len alone | EAGLE bigram/page alignment can make final insertion find a longer existing prefix and force duplicate rollback work. Confidence: medium Scope-risk: moderate Directive: Do not add per-layer CP collectives here; keep radix state synchronous and data transfer asynchronous/local-event driven. Tested: local git diff --check Tested: local py_compile for touched CP HiCache/cache-controller/deepseek/test files Tested: remote pytest test/registered/unit/mem_cache/test_cp_hicache_metadata.py test/registered/unit/managers/test_hicache_controller_cp.py -q (115 passed, 5 warnings) Not-tested: full GLM5 ETE server rerun after this commit
2324 lines
88 KiB
Python
2324 lines
88 KiB
Python
# Copyright 2023-2024 SGLang Team
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
# ==============================================================================
|
|
|
|
# Adapted from:
|
|
# https://github.com/vllm-project/vllm/blob/fb6af8bc086328ca6659e72d11ffd4309ce4de22/vllm/model_executor/models/deepseek_v2.py
|
|
"""Inference-only DeepseekV2 model."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from contextlib import nullcontext
|
|
from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from torch import nn
|
|
from transformers import PretrainedConfig
|
|
|
|
from sglang.srt.batch_overlap.single_batch_overlap import SboFlags, compute_overlap_args
|
|
from sglang.srt.batch_overlap.two_batch_overlap import (
|
|
MaybeTboDeepEPDispatcher,
|
|
model_forward_maybe_tbo,
|
|
)
|
|
from sglang.srt.configs.model_config import (
|
|
compute_mla_mscale_scaling,
|
|
get_nsa_index_head_dim,
|
|
get_nsa_index_n_heads,
|
|
get_nsa_index_topk,
|
|
is_deepseek_nsa,
|
|
)
|
|
from sglang.srt.distributed import (
|
|
divide,
|
|
get_moe_expert_parallel_world_size,
|
|
get_pp_group,
|
|
get_tensor_model_parallel_world_size,
|
|
tensor_model_parallel_all_reduce,
|
|
)
|
|
from sglang.srt.environ import envs
|
|
from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
|
|
from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation
|
|
from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo
|
|
from sglang.srt.layers import deep_gemm_wrapper
|
|
from sglang.srt.layers.activation import SiluAndMul
|
|
from sglang.srt.layers.amx_utils import PackWeightMethod
|
|
from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer
|
|
from sglang.srt.layers.attention.nsa.utils import (
|
|
can_cp_split,
|
|
cp_all_gather_rerange_output,
|
|
cp_collect_last_token_hidden,
|
|
cp_split_and_rebuild_data,
|
|
cp_split_and_rebuild_position,
|
|
is_nsa_enable_prefill_cp,
|
|
nsa_use_prefill_cp,
|
|
prepare_input_dp_with_cp_dsa,
|
|
)
|
|
from sglang.srt.layers.communicator import (
|
|
LayerCommunicator,
|
|
LayerScatterModes,
|
|
enable_moe_dense_fully_dp,
|
|
get_attn_tp_context,
|
|
)
|
|
from sglang.srt.layers.communicator_nsa_cp import NSACPLayerCommunicator
|
|
from sglang.srt.layers.dp_attention import (
|
|
get_attention_cp_rank,
|
|
get_attention_cp_size,
|
|
get_attention_tp_group,
|
|
get_attention_tp_rank,
|
|
get_attention_tp_size,
|
|
is_dp_attention_enabled,
|
|
)
|
|
from sglang.srt.layers.layernorm import RMSNorm
|
|
from sglang.srt.layers.linear import (
|
|
ColumnParallelLinear,
|
|
MergedColumnParallelLinear,
|
|
ReplicatedLinear,
|
|
RowParallelLinear,
|
|
)
|
|
from sglang.srt.layers.logits_processor import LogitsProcessor
|
|
from sglang.srt.layers.moe import (
|
|
get_moe_a2a_backend,
|
|
get_moe_runner_backend,
|
|
should_use_flashinfer_cutlass_moe_fp4_allgather,
|
|
)
|
|
from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class
|
|
from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
|
|
from sglang.srt.layers.moe.kt_ep_wrapper import KTEPWrapperMethod
|
|
from sglang.srt.layers.moe.token_dispatcher.base import (
|
|
BaseDispatcher,
|
|
CombineInput,
|
|
DispatchOutput,
|
|
)
|
|
from sglang.srt.layers.moe.topk import TopK, TopKOutputFormat
|
|
from sglang.srt.layers.moe.utils import (
|
|
RoutingMethodType,
|
|
filter_moe_weight_param_global_expert,
|
|
)
|
|
from sglang.srt.layers.quantization.base_config import QuantizationConfig
|
|
from sglang.srt.layers.quantization.fp8 import Fp8Config
|
|
from sglang.srt.layers.radix_attention import RadixAttention
|
|
from sglang.srt.layers.rotary_embedding import get_rope_wrapper
|
|
from sglang.srt.layers.utils import PPMissingLayer
|
|
from sglang.srt.layers.vocab_parallel_embedding import (
|
|
ParallelLMHead,
|
|
VocabParallelEmbedding,
|
|
)
|
|
from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
|
|
from sglang.srt.model_executor.forward_batch_info import (
|
|
CaptureHiddenMode,
|
|
ForwardBatch,
|
|
PPProxyTensors,
|
|
)
|
|
from sglang.srt.models.deepseek_common.attention_backend_handler import (
|
|
AttentionBackendRegistry,
|
|
)
|
|
from sglang.srt.models.deepseek_common.attention_forward_methods import (
|
|
AttnForwardMethod,
|
|
DeepseekMHAForwardMixin,
|
|
DeepseekMLACpuForwardMixin,
|
|
DeepseekMLAForwardMixin,
|
|
DeepseekMLARocmForwardMixin,
|
|
)
|
|
from sglang.srt.models.deepseek_common.deepseek_weight_loader import (
|
|
DeepseekV2WeightLoaderMixin,
|
|
)
|
|
from sglang.srt.models.deepseek_common.utils import (
|
|
_device_sm,
|
|
_get_llama_4_scaling,
|
|
_is_cpu,
|
|
_is_cpu_amx_available,
|
|
_is_cuda,
|
|
_is_gfx95_supported,
|
|
_is_hip,
|
|
_is_npu,
|
|
_use_aiter,
|
|
_use_aiter_gfx95,
|
|
)
|
|
from sglang.srt.server_args import get_global_server_args
|
|
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
|
|
from sglang.srt.utils import (
|
|
BumpAllocator,
|
|
LazyValue,
|
|
add_prefix,
|
|
is_non_idle_and_non_empty,
|
|
log_info_on_rank0,
|
|
make_layers,
|
|
use_intel_amx_backend,
|
|
)
|
|
|
|
if _use_aiter_gfx95:
|
|
from sglang.srt.layers.rocm_linear_utils import (
|
|
aiter_dsv3_router_gemm,
|
|
get_dsv3_gemm_output_zero_allocator_size,
|
|
)
|
|
|
|
if _use_aiter:
|
|
pass
|
|
|
|
if _is_cuda:
|
|
from sgl_kernel import dsv3_fused_a_gemm, dsv3_router_gemm
|
|
elif _is_npu:
|
|
from sglang.srt.hardware_backend.npu.modules.deepseek_v2_attention_mla_npu import (
|
|
forward_dsa_core_npu,
|
|
forward_dsa_prepare_npu,
|
|
forward_mha_core_npu,
|
|
forward_mha_prepare_npu,
|
|
forward_mla_core_npu,
|
|
forward_mla_prepare_npu,
|
|
)
|
|
else:
|
|
pass
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class DeepseekV2MLP(nn.Module):
|
|
def __init__(
|
|
self,
|
|
hidden_size: int,
|
|
intermediate_size: int,
|
|
hidden_act: str,
|
|
quant_config: Optional[QuantizationConfig] = None,
|
|
reduce_results: bool = True,
|
|
prefix: str = "",
|
|
tp_rank: Optional[int] = None,
|
|
tp_size: Optional[int] = None,
|
|
) -> None:
|
|
super().__init__()
|
|
self.tp_size = tp_size
|
|
|
|
self.gate_up_proj = MergedColumnParallelLinear(
|
|
hidden_size,
|
|
[intermediate_size] * 2,
|
|
bias=False,
|
|
quant_config=quant_config,
|
|
prefix=add_prefix("gate_up_proj", prefix),
|
|
tp_rank=tp_rank,
|
|
tp_size=tp_size,
|
|
)
|
|
self.down_proj = RowParallelLinear(
|
|
intermediate_size,
|
|
hidden_size,
|
|
bias=False,
|
|
quant_config=quant_config,
|
|
reduce_results=reduce_results,
|
|
prefix=add_prefix("down_proj", prefix),
|
|
tp_rank=tp_rank,
|
|
tp_size=tp_size,
|
|
)
|
|
if not hasattr(self.gate_up_proj, "weight") and hasattr(
|
|
self.gate_up_proj, "weight_packed"
|
|
):
|
|
self.gate_up_proj.weight = self.gate_up_proj.weight_packed
|
|
if not hasattr(self.down_proj, "weight") and hasattr(
|
|
self.down_proj, "weight_packed"
|
|
):
|
|
self.down_proj.weight = self.down_proj.weight_packed
|
|
if hidden_act != "silu":
|
|
raise ValueError(
|
|
f"Unsupported activation: {hidden_act}. Only silu is supported for now."
|
|
)
|
|
self.act_fn = SiluAndMul()
|
|
|
|
def forward(
|
|
self,
|
|
x,
|
|
forward_batch=None,
|
|
should_allreduce_fusion: bool = False,
|
|
use_reduce_scatter: bool = False,
|
|
gemm_output_zero_allocator: BumpAllocator = None,
|
|
):
|
|
if (self.tp_size == 1) and x.shape[0] == 0:
|
|
return x
|
|
|
|
if (
|
|
gemm_output_zero_allocator is not None
|
|
and x.shape[0] <= 256
|
|
and self.gate_up_proj.weight.dtype == torch.uint8
|
|
):
|
|
y = gemm_output_zero_allocator.allocate(
|
|
x.shape[0] * self.gate_up_proj.output_size_per_partition
|
|
).view(x.shape[0], self.gate_up_proj.output_size_per_partition)
|
|
x = (x, None, y)
|
|
|
|
gate_up, _ = self.gate_up_proj(x)
|
|
x = self.act_fn(gate_up)
|
|
x, _ = self.down_proj(
|
|
x,
|
|
skip_all_reduce=should_allreduce_fusion or use_reduce_scatter,
|
|
)
|
|
return x
|
|
|
|
|
|
class MoEGate(nn.Module):
|
|
def __init__(
|
|
self,
|
|
config,
|
|
quant_config,
|
|
prefix: str = "",
|
|
is_nextn: bool = False,
|
|
):
|
|
super().__init__()
|
|
self.is_nextn = is_nextn
|
|
self.weight = nn.Parameter(
|
|
torch.empty((config.n_routed_experts, config.hidden_size))
|
|
)
|
|
if config.topk_method == "noaux_tc":
|
|
correction_bias_dtype = torch.float32
|
|
if quant_config is not None:
|
|
if (
|
|
quant_config.get_name() == "modelopt_fp4"
|
|
and get_moe_runner_backend().is_flashinfer_trtllm()
|
|
):
|
|
correction_bias_dtype = torch.bfloat16
|
|
elif _use_aiter and quant_config.get_name() in (
|
|
"fp8",
|
|
"compressed_tensors",
|
|
"quark",
|
|
):
|
|
correction_bias_dtype = torch.bfloat16
|
|
self.e_score_correction_bias = nn.Parameter(
|
|
torch.empty((config.n_routed_experts), dtype=correction_bias_dtype)
|
|
)
|
|
else:
|
|
self.e_score_correction_bias = None
|
|
if _is_cpu and _is_cpu_amx_available:
|
|
self.quant_method = PackWeightMethod(weight_names=["weight"])
|
|
self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp()
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states,
|
|
gemm_output_zero_allocator: BumpAllocator = None,
|
|
forward_batch: ForwardBatch = None,
|
|
):
|
|
if use_intel_amx_backend(self):
|
|
return torch.ops.sgl_kernel.weight_packed_linear(
|
|
hidden_states,
|
|
self.weight,
|
|
None, # bias
|
|
True, # is_vnni
|
|
)
|
|
|
|
if get_global_server_args().enable_deterministic_inference:
|
|
return F.linear(hidden_states, self.weight, None)
|
|
|
|
if forward_batch is not None and nsa_use_prefill_cp(forward_batch):
|
|
logits = F.linear(hidden_states, self.weight, None)
|
|
else:
|
|
# NOTE: For some unknown reason, router_gemm seems degrade accept length.
|
|
if (
|
|
_is_cuda
|
|
and hidden_states.shape[0] <= 16
|
|
and hidden_states.shape[1] == 7168
|
|
and (self.weight.shape[0] == 256 or self.weight.shape[0] == 384)
|
|
and _device_sm >= 90
|
|
):
|
|
# router gemm output float32
|
|
logits = dsv3_router_gemm(
|
|
hidden_states, self.weight, out_dtype=torch.float32
|
|
)
|
|
elif (
|
|
_use_aiter_gfx95
|
|
and hidden_states.shape[0] <= 256
|
|
and self.weight.shape[0] <= 256
|
|
):
|
|
logits = aiter_dsv3_router_gemm(
|
|
hidden_states, self.weight, gemm_output_zero_allocator
|
|
)
|
|
else:
|
|
logits = F.linear(hidden_states, self.weight, None)
|
|
|
|
return logits
|
|
|
|
|
|
class DeepseekV2MoE(nn.Module):
|
|
def __init__(
|
|
self,
|
|
config: PretrainedConfig,
|
|
layer_id: int,
|
|
quant_config: Optional[QuantizationConfig] = None,
|
|
prefix: str = "",
|
|
alt_stream: Optional[torch.cuda.Stream] = None,
|
|
is_nextn: bool = False,
|
|
):
|
|
super().__init__()
|
|
self.tp_size = get_tensor_model_parallel_world_size()
|
|
self.moe_ep_size = get_moe_expert_parallel_world_size()
|
|
self.routed_scaling_factor = config.routed_scaling_factor
|
|
self.n_shared_experts = config.n_shared_experts
|
|
self.num_fused_shared_experts = (
|
|
0
|
|
if get_global_server_args().disable_shared_experts_fusion
|
|
else config.n_shared_experts
|
|
)
|
|
self.config = config
|
|
self.layer_id = layer_id
|
|
self.alt_stream = alt_stream
|
|
self.is_nextn = is_nextn
|
|
|
|
if self.tp_size > config.n_routed_experts:
|
|
raise ValueError(
|
|
f"Tensor parallel size {self.tp_size} is greater than "
|
|
f"the number of experts {config.n_routed_experts}."
|
|
)
|
|
|
|
if config.hidden_act != "silu":
|
|
raise ValueError(
|
|
f"Unsupported activation: {config.hidden_act}. "
|
|
"Only silu is supported for now."
|
|
)
|
|
|
|
self.gate = MoEGate(
|
|
config=config,
|
|
quant_config=quant_config,
|
|
prefix=add_prefix("gate", prefix),
|
|
is_nextn=is_nextn,
|
|
)
|
|
|
|
# scaling factor for fused shared experts on AMD-platform.
|
|
fused_shared_experts_scaling_factor = None
|
|
if self.moe_ep_size > 1 and self.num_fused_shared_experts > 0:
|
|
# if enable_ep_moe tp_szie == ep_size, every gpu get shared experts gemm output
|
|
# so we scale with 1 / self.moe_ep_size in ep mode which will make it equalation as in tp mode
|
|
# with fused_shared_experts
|
|
fused_shared_experts_scaling_factor = 1.0 / float(self.moe_ep_size)
|
|
|
|
self.experts = get_moe_impl_class(quant_config)(
|
|
num_experts=config.n_routed_experts
|
|
+ self.num_fused_shared_experts
|
|
+ get_global_server_args().ep_num_redundant_experts,
|
|
num_fused_shared_experts=self.num_fused_shared_experts,
|
|
top_k=config.num_experts_per_tok + self.num_fused_shared_experts,
|
|
hidden_size=config.hidden_size,
|
|
intermediate_size=config.moe_intermediate_size,
|
|
layer_id=self.layer_id,
|
|
quant_config=quant_config,
|
|
routed_scaling_factor=self.routed_scaling_factor,
|
|
routing_method_type=getattr(
|
|
config, "routing_method_type", RoutingMethodType.DeepSeekV3
|
|
),
|
|
prefix=add_prefix("experts", prefix),
|
|
)
|
|
|
|
self.topk = TopK(
|
|
top_k=config.num_experts_per_tok + self.num_fused_shared_experts,
|
|
layer_id=self.layer_id,
|
|
renormalize=config.norm_topk_prob,
|
|
use_grouped_topk=True,
|
|
num_expert_group=config.n_group,
|
|
num_fused_shared_experts=self.num_fused_shared_experts,
|
|
topk_group=config.topk_group,
|
|
correction_bias=self.gate.e_score_correction_bias,
|
|
quant_config=quant_config,
|
|
routed_scaling_factor=self.routed_scaling_factor,
|
|
apply_routed_scaling_factor_on_output=self.experts.should_fuse_routed_scaling_factor_in_topk,
|
|
fused_shared_experts_scaling_factor=fused_shared_experts_scaling_factor,
|
|
# Some Fp4 MoE backends require the output format to be bypassed but the MTP layers are unquantized
|
|
# and requires the output format to be standard (except trtllm). We use quant_config to determine the output format.
|
|
output_format=(
|
|
TopKOutputFormat.STANDARD
|
|
if (quant_config is None)
|
|
and (not get_moe_runner_backend().is_flashinfer_trtllm())
|
|
else None
|
|
),
|
|
)
|
|
|
|
self.shared_experts_is_int8 = False
|
|
self.shared_experts_is_fp8 = False
|
|
self.shared_experts_weight_block_size = None
|
|
if config.n_shared_experts is not None and self.num_fused_shared_experts == 0:
|
|
intermediate_size = config.moe_intermediate_size * config.n_shared_experts
|
|
# disable tp for shared experts when enable deepep moe, or with fp4 allgather
|
|
self.shared_experts = DeepseekV2MLP(
|
|
hidden_size=config.hidden_size,
|
|
intermediate_size=intermediate_size,
|
|
hidden_act=config.hidden_act,
|
|
quant_config=quant_config,
|
|
reduce_results=False,
|
|
prefix=add_prefix("shared_experts", prefix),
|
|
**(
|
|
dict(tp_rank=0, tp_size=1)
|
|
if get_moe_a2a_backend().is_deepep()
|
|
or get_moe_a2a_backend().is_mooncake()
|
|
or get_moe_a2a_backend().is_nixl()
|
|
or get_moe_a2a_backend().is_mori()
|
|
or get_moe_a2a_backend().is_ascend_fuseep()
|
|
or get_moe_a2a_backend().is_flashinfer()
|
|
or should_use_flashinfer_cutlass_moe_fp4_allgather()
|
|
else {}
|
|
),
|
|
)
|
|
is_packed_weight = (
|
|
hasattr(self.shared_experts.gate_up_proj.quant_method, "quant_config")
|
|
and self.shared_experts.gate_up_proj.quant_method.quant_config.get_name()
|
|
in {
|
|
"awq",
|
|
"awq_marlin",
|
|
"moe_wna16",
|
|
}
|
|
)
|
|
self.shared_experts_is_int8 = (
|
|
not is_packed_weight
|
|
and self.shared_experts.gate_up_proj.weight.dtype == torch.int8
|
|
)
|
|
self.shared_experts_is_fp8 = (
|
|
not is_packed_weight
|
|
and self.shared_experts.gate_up_proj.weight.dtype == torch.float8_e4m3fn
|
|
)
|
|
if self.shared_experts_is_fp8:
|
|
if (
|
|
_use_aiter
|
|
and config.quantization_config.get("quant_method")
|
|
== "compressed-tensors"
|
|
):
|
|
# For compressed-tensors ptpc model, don't need to check the weight_block_size
|
|
pass
|
|
else:
|
|
assert (
|
|
self.shared_experts.gate_up_proj.quant_method.quant_config.weight_block_size
|
|
== self.shared_experts.down_proj.quant_method.quant_config.weight_block_size
|
|
)
|
|
self.shared_experts_weight_block_size = self.shared_experts.gate_up_proj.quant_method.quant_config.weight_block_size
|
|
|
|
self.top_k = config.num_experts_per_tok
|
|
|
|
if (
|
|
get_moe_a2a_backend().is_deepep()
|
|
or get_moe_a2a_backend().is_mooncake()
|
|
or get_moe_a2a_backend().is_nixl()
|
|
or get_moe_a2a_backend().is_mori()
|
|
or get_moe_a2a_backend().is_ascend_fuseep()
|
|
):
|
|
# TODO: we will support tp < ep in the future
|
|
self.ep_size = get_moe_expert_parallel_world_size()
|
|
self.num_experts = (
|
|
config.n_routed_experts
|
|
+ get_global_server_args().ep_num_redundant_experts
|
|
)
|
|
self.renormalize = config.norm_topk_prob
|
|
self.topk_group = config.topk_group
|
|
self.num_expert_group = config.n_group
|
|
self.correction_bias = (
|
|
self.gate.e_score_correction_bias.data
|
|
if self.gate.e_score_correction_bias is not None
|
|
else None
|
|
)
|
|
|
|
self._enable_a2a_moe = (
|
|
get_moe_a2a_backend().is_deepep()
|
|
or get_moe_a2a_backend().is_mooncake()
|
|
or get_moe_a2a_backend().is_nixl()
|
|
or get_moe_a2a_backend().is_mori()
|
|
or get_moe_a2a_backend().is_ascend_fuseep()
|
|
or get_moe_a2a_backend().is_flashinfer()
|
|
)
|
|
self._fuse_shared_experts_inside_sbo = SboFlags.fuse_shared_experts_inside_sbo()
|
|
|
|
def get_moe_weights(self):
|
|
return [
|
|
x.data
|
|
for name, x in self.experts.named_parameters()
|
|
if name not in ["correction_bias"]
|
|
and filter_moe_weight_param_global_expert(
|
|
name, x, self.experts.num_local_experts
|
|
)
|
|
]
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
forward_batch: Optional[ForwardBatch] = None,
|
|
should_allreduce_fusion: bool = False,
|
|
use_reduce_scatter: bool = False,
|
|
gemm_output_zero_allocator: BumpAllocator = None,
|
|
) -> torch.Tensor:
|
|
if not self._enable_a2a_moe:
|
|
if (
|
|
self.alt_stream is not None
|
|
and self.num_fused_shared_experts == 0
|
|
and hidden_states.shape[0] > 0
|
|
and get_is_capture_mode()
|
|
):
|
|
return self.forward_normal_dual_stream(
|
|
hidden_states,
|
|
should_allreduce_fusion,
|
|
use_reduce_scatter,
|
|
gemm_output_zero_allocator,
|
|
)
|
|
else:
|
|
return self.forward_normal(
|
|
hidden_states,
|
|
should_allreduce_fusion,
|
|
use_reduce_scatter,
|
|
gemm_output_zero_allocator,
|
|
)
|
|
else:
|
|
return self.forward_deepep(hidden_states, forward_batch)
|
|
|
|
def forward_normal_dual_stream(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
should_allreduce_fusion: bool = False,
|
|
use_reduce_scatter: bool = False,
|
|
gemm_output_zero_allocator: BumpAllocator = None,
|
|
) -> torch.Tensor:
|
|
current_stream = torch.cuda.current_stream()
|
|
self.alt_stream.wait_stream(current_stream)
|
|
shared_output = self._forward_shared_experts(
|
|
hidden_states, gemm_output_zero_allocator
|
|
)
|
|
server_args = get_global_server_args()
|
|
dispatch_info = (
|
|
ExpertLocationDispatchInfo.init_new(layer_id=self.layer_id)
|
|
if server_args.enable_eplb
|
|
else None
|
|
)
|
|
with torch.cuda.stream(self.alt_stream):
|
|
# router_logits: (num_tokens, n_experts)
|
|
router_logits = self.gate(hidden_states, gemm_output_zero_allocator)
|
|
topk_output = self.topk(
|
|
hidden_states,
|
|
router_logits,
|
|
expert_location_dispatch_info=dispatch_info,
|
|
)
|
|
final_hidden_states = self.experts(hidden_states, topk_output)
|
|
if not _is_cuda or isinstance(self.experts.quant_method, KTEPWrapperMethod):
|
|
final_hidden_states *= self.routed_scaling_factor
|
|
|
|
current_stream.wait_stream(self.alt_stream)
|
|
final_hidden_states += shared_output
|
|
if (
|
|
self.tp_size > 1
|
|
and not should_allreduce_fusion
|
|
and not use_reduce_scatter
|
|
and not should_use_flashinfer_cutlass_moe_fp4_allgather()
|
|
):
|
|
final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
|
|
return final_hidden_states
|
|
|
|
def forward_normal(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
should_allreduce_fusion: bool = False,
|
|
use_reduce_scatter: bool = False,
|
|
gemm_output_zero_allocator: BumpAllocator = None,
|
|
) -> torch.Tensor:
|
|
if hasattr(self, "shared_experts") and use_intel_amx_backend(
|
|
self.shared_experts.gate_up_proj
|
|
):
|
|
return self.forward_cpu(hidden_states, should_allreduce_fusion)
|
|
server_args = get_global_server_args()
|
|
dispatch_info = (
|
|
ExpertLocationDispatchInfo.init_new(layer_id=self.layer_id)
|
|
if server_args.enable_eplb
|
|
else None
|
|
)
|
|
if hidden_states.shape[0] > 0:
|
|
if (
|
|
not self._fuse_shared_experts_inside_sbo
|
|
): # TODO: check if it supports mtp
|
|
shared_output = self._forward_shared_experts(
|
|
hidden_states, gemm_output_zero_allocator
|
|
)
|
|
# router_logits: (num_tokens, n_experts)
|
|
router_logits = self.gate(hidden_states, gemm_output_zero_allocator)
|
|
topk_output = self.topk(
|
|
hidden_states,
|
|
router_logits,
|
|
expert_location_dispatch_info=dispatch_info,
|
|
)
|
|
else:
|
|
shared_output = None
|
|
topk_output = self.topk.empty_topk_output(hidden_states.device)
|
|
|
|
if self._fuse_shared_experts_inside_sbo:
|
|
shared_output = None
|
|
|
|
def _pre_combine_hook(
|
|
dispatcher: BaseDispatcher, combine_input: CombineInput
|
|
):
|
|
nonlocal shared_output
|
|
self.alt_stream.wait_stream(torch.cuda.current_stream())
|
|
with torch.cuda.stream(self.alt_stream):
|
|
shared_output = self._forward_shared_experts(
|
|
hidden_states, gemm_output_zero_allocator
|
|
)
|
|
|
|
pre_combine_hook_handle.remove()
|
|
|
|
def _post_combine_hook(
|
|
dispatcher: BaseDispatcher, hidden_states: torch.Tensor
|
|
):
|
|
nonlocal shared_output
|
|
torch.cuda.current_stream().wait_stream(self.alt_stream)
|
|
post_combine_hook_handle.remove()
|
|
|
|
pre_combine_hook_handle = self.experts.dispatcher.register_pre_combine_hook(
|
|
_pre_combine_hook
|
|
)
|
|
post_combine_hook_handle = (
|
|
self.experts.dispatcher.register_post_combine_hook(_post_combine_hook)
|
|
)
|
|
|
|
final_hidden_states = self.experts(
|
|
hidden_states,
|
|
topk_output,
|
|
)
|
|
if (
|
|
not _is_cuda
|
|
and not _use_aiter
|
|
or isinstance(self.experts.quant_method, KTEPWrapperMethod)
|
|
):
|
|
# fused in biased_grouped_topk so we can skip here
|
|
final_hidden_states *= self.routed_scaling_factor
|
|
if shared_output is not None:
|
|
final_hidden_states += shared_output
|
|
if (
|
|
self.tp_size > 1
|
|
and not should_allreduce_fusion
|
|
and not use_reduce_scatter
|
|
and not should_use_flashinfer_cutlass_moe_fp4_allgather()
|
|
):
|
|
final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
|
|
return final_hidden_states
|
|
|
|
def forward_cpu(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
should_allreduce_fusion: bool = False,
|
|
) -> torch.Tensor:
|
|
# router_logits: (num_tokens, n_experts)
|
|
router_logits = self.gate(hidden_states)
|
|
topk_output = self.topk(hidden_states, router_logits)
|
|
fused_experts_out = self.experts(
|
|
hidden_states=hidden_states, topk_output=topk_output
|
|
)
|
|
|
|
assert use_intel_amx_backend(
|
|
self.shared_experts.gate_up_proj
|
|
) == use_intel_amx_backend(self.shared_experts.down_proj)
|
|
# [Note] inplace should be False in fused_experts.
|
|
# If inplace is True in fused_experts (self.experts), hidden_states will be changed after fused_experts
|
|
# While hidden_states is still needed in shared_expert.
|
|
final_hidden_states = torch.ops.sgl_kernel.shared_expert_cpu(
|
|
hidden_states,
|
|
self.shared_experts.gate_up_proj.weight,
|
|
self.shared_experts.down_proj.weight,
|
|
fused_experts_out,
|
|
self.routed_scaling_factor,
|
|
True, # inplace
|
|
self.shared_experts_is_int8, # use_int8_w8a8
|
|
self.shared_experts_is_fp8, # use_fp8_w8a16
|
|
(
|
|
self.shared_experts.gate_up_proj.weight_scale
|
|
if self.shared_experts_is_int8
|
|
else (
|
|
self.shared_experts.gate_up_proj.weight_scale_inv
|
|
if self.shared_experts_is_fp8
|
|
else None
|
|
)
|
|
), # w1_scale
|
|
(
|
|
self.shared_experts.down_proj.weight_scale
|
|
if self.shared_experts_is_int8
|
|
else (
|
|
self.shared_experts.down_proj.weight_scale_inv
|
|
if self.shared_experts_is_fp8
|
|
else None
|
|
)
|
|
), # w2_scale
|
|
(
|
|
self.shared_experts_weight_block_size
|
|
if self.shared_experts_is_fp8
|
|
else None
|
|
), # block_size
|
|
True, # is_vnni
|
|
)
|
|
if self.tp_size > 1 and not should_allreduce_fusion:
|
|
final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
|
|
return final_hidden_states
|
|
|
|
def forward_deepep(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
forward_batch: ForwardBatch,
|
|
) -> torch.Tensor:
|
|
shared_output = None
|
|
sbo_enabled_flag = self._fuse_shared_experts_inside_sbo and not self.is_nextn
|
|
sbo_overlap_dispatch_flag = (
|
|
sbo_enabled_flag and SboFlags.enable_dispatch_shared_one_stream_overlap()
|
|
)
|
|
sbo_overlap_combine_flag = (
|
|
sbo_enabled_flag and SboFlags.enable_combine_shared_two_stream_overlap()
|
|
)
|
|
|
|
if hidden_states.shape[0] > 0:
|
|
# router_logits: (num_tokens, n_experts)
|
|
router_logits = self.gate(hidden_states, forward_batch=forward_batch)
|
|
if not sbo_enabled_flag:
|
|
if self.alt_stream is not None:
|
|
self.alt_stream.wait_stream(torch.cuda.current_stream())
|
|
with torch.cuda.stream(self.alt_stream):
|
|
shared_output = self._forward_shared_experts(hidden_states)
|
|
shared_output.record_stream(self.alt_stream)
|
|
shared_event = self.alt_stream.record_event()
|
|
else:
|
|
shared_output = self._forward_shared_experts(hidden_states)
|
|
topk_output = self.topk(
|
|
hidden_states,
|
|
router_logits,
|
|
num_token_non_padded=forward_batch.num_token_non_padded,
|
|
expert_location_dispatch_info=ExpertLocationDispatchInfo.init_new(
|
|
layer_id=self.layer_id,
|
|
),
|
|
)
|
|
else:
|
|
topk_output = self.topk.empty_topk_output(hidden_states.device)
|
|
|
|
if sbo_overlap_dispatch_flag:
|
|
shared_output = None
|
|
|
|
def _deepep_dispatch_hook(dispatcher: BaseDispatcher):
|
|
nonlocal shared_output
|
|
shared_output = self._forward_shared_experts(hidden_states)
|
|
for handle in deepep_dispatch_hook_handle:
|
|
handle.remove()
|
|
|
|
def _post_dispatch_hook(
|
|
dispatcher: BaseDispatcher, dispatch_output: DispatchOutput
|
|
):
|
|
combine_overlap_args, down_gemm_overlap_args, meta_overlap_args = (
|
|
compute_overlap_args(dispatch_output, self.alt_stream)
|
|
)
|
|
dispatcher.set_overlap_args(
|
|
combine_overlap_args=combine_overlap_args,
|
|
meta_overlap_args=meta_overlap_args,
|
|
)
|
|
self.experts.set_overlap_args(
|
|
down_gemm_overlap_args=down_gemm_overlap_args,
|
|
meta_overlap_args=meta_overlap_args,
|
|
)
|
|
post_dispatch_hook_handle.remove()
|
|
|
|
def _post_combine_hook(
|
|
dispatcher: BaseDispatcher, hidden_states: torch.Tensor
|
|
):
|
|
dispatcher.clear_overlap_args()
|
|
self.experts.clear_overlap_args()
|
|
post_combine_hook_handle.remove()
|
|
|
|
assert isinstance(self.experts.dispatcher, MaybeTboDeepEPDispatcher)
|
|
deepep_dispatch_hook_handle = (
|
|
self.experts.dispatcher.register_deepep_dispatch_hook(
|
|
_deepep_dispatch_hook
|
|
)
|
|
)
|
|
post_dispatch_hook_handle = (
|
|
self.experts.dispatcher.register_post_dispatch_hook(_post_dispatch_hook)
|
|
)
|
|
post_combine_hook_handle = (
|
|
self.experts.dispatcher.register_post_combine_hook(_post_combine_hook)
|
|
)
|
|
|
|
elif sbo_overlap_combine_flag:
|
|
shared_output = None
|
|
|
|
def _post_dispatch_hook(
|
|
dispatcher: BaseDispatcher, dispatch_output: DispatchOutput
|
|
):
|
|
combine_overlap_args, down_gemm_overlap_args, meta_overlap_args = (
|
|
compute_overlap_args(dispatch_output, self.alt_stream)
|
|
)
|
|
dispatcher.set_overlap_args(
|
|
combine_overlap_args=combine_overlap_args,
|
|
meta_overlap_args=meta_overlap_args,
|
|
)
|
|
self.experts.set_overlap_args(
|
|
down_gemm_overlap_args=down_gemm_overlap_args,
|
|
meta_overlap_args=meta_overlap_args,
|
|
)
|
|
|
|
post_dispatch_hook_handle.remove()
|
|
|
|
def _pre_combine_hook(
|
|
dispatcher: BaseDispatcher, combine_input: CombineInput
|
|
):
|
|
nonlocal shared_output
|
|
|
|
if (
|
|
e := dispatcher.meta_overlap_args.get("record_event_after_down")
|
|
) is not None:
|
|
e.record()
|
|
|
|
# TODO reduce sm for non-deepgemm
|
|
with deep_gemm_wrapper.configure_deep_gemm_num_sms(
|
|
dispatcher.meta_overlap_args["compute_num_sms"]
|
|
):
|
|
shared_output = self._forward_shared_experts(hidden_states)
|
|
|
|
pre_combine_hook_handle.remove()
|
|
|
|
def _post_combine_hook(
|
|
dispatcher: BaseDispatcher, hidden_states: torch.Tensor
|
|
):
|
|
dispatcher.clear_overlap_args()
|
|
self.experts.clear_overlap_args()
|
|
post_combine_hook_handle.remove()
|
|
|
|
post_dispatch_hook_handle = (
|
|
self.experts.dispatcher.register_post_dispatch_hook(_post_dispatch_hook)
|
|
)
|
|
pre_combine_hook_handle = self.experts.dispatcher.register_pre_combine_hook(
|
|
_pre_combine_hook
|
|
)
|
|
post_combine_hook_handle = (
|
|
self.experts.dispatcher.register_post_combine_hook(_post_combine_hook)
|
|
)
|
|
elif envs.SGLANG_BLACKWELL_OVERLAP_SHARED_EXPERTS_OUTSIDE_SBO.get():
|
|
# On GB200: Shared experts overlapped on alt_stream, down gemm overlapped with DeepEP Combine
|
|
|
|
def _post_dispatch_hook(
|
|
dispatcher: BaseDispatcher, dispatch_output: DispatchOutput
|
|
):
|
|
combine_overlap_args, down_gemm_overlap_args, meta_overlap_args = (
|
|
compute_overlap_args(dispatch_output, self.alt_stream)
|
|
)
|
|
dispatcher.set_overlap_args(
|
|
combine_overlap_args=combine_overlap_args,
|
|
meta_overlap_args=meta_overlap_args,
|
|
)
|
|
self.experts.set_overlap_args(
|
|
down_gemm_overlap_args=down_gemm_overlap_args,
|
|
meta_overlap_args=meta_overlap_args,
|
|
)
|
|
|
|
post_dispatch_hook_handle.remove()
|
|
|
|
def _pre_combine_hook(
|
|
dispatcher: BaseDispatcher, combine_input: CombineInput
|
|
):
|
|
if (
|
|
e := dispatcher.meta_overlap_args.get("record_event_after_down")
|
|
) is not None:
|
|
e.record()
|
|
pre_combine_hook_handle.remove()
|
|
|
|
def _post_combine_hook(
|
|
dispatcher: BaseDispatcher, hidden_states: torch.Tensor
|
|
):
|
|
dispatcher.clear_overlap_args()
|
|
self.experts.clear_overlap_args()
|
|
post_combine_hook_handle.remove()
|
|
|
|
post_dispatch_hook_handle = (
|
|
self.experts.dispatcher.register_post_dispatch_hook(_post_dispatch_hook)
|
|
)
|
|
pre_combine_hook_handle = self.experts.dispatcher.register_pre_combine_hook(
|
|
_pre_combine_hook
|
|
)
|
|
post_combine_hook_handle = (
|
|
self.experts.dispatcher.register_post_combine_hook(_post_combine_hook)
|
|
)
|
|
|
|
final_hidden_states = self.experts(
|
|
hidden_states=hidden_states,
|
|
topk_output=topk_output,
|
|
)
|
|
|
|
if (
|
|
hidden_states.shape[0] > 0
|
|
and not sbo_enabled_flag
|
|
and self.alt_stream is not None
|
|
):
|
|
torch.cuda.current_stream().wait_event(shared_event)
|
|
|
|
if shared_output is not None:
|
|
x = shared_output
|
|
# aiter moe call will handle routed_scaling_factor in the function
|
|
# so add _use_aiter condition to eliminate to use self.routed_scaling_factor in add_ call
|
|
if self.experts.should_fuse_routed_scaling_factor_in_topk or _use_aiter:
|
|
x.add_(final_hidden_states)
|
|
else:
|
|
x.add_(final_hidden_states, alpha=self.routed_scaling_factor)
|
|
final_hidden_states = x
|
|
else:
|
|
if not (
|
|
self.experts.should_fuse_routed_scaling_factor_in_topk or _use_aiter
|
|
):
|
|
final_hidden_states *= self.routed_scaling_factor
|
|
|
|
return final_hidden_states
|
|
|
|
def _forward_shared_experts(
|
|
self, hidden_states, gemm_output_zero_allocator: BumpAllocator = None
|
|
):
|
|
if (hidden_states.shape[0] > 0) and (self.num_fused_shared_experts == 0):
|
|
return self.shared_experts(
|
|
hidden_states, gemm_output_zero_allocator=gemm_output_zero_allocator
|
|
)
|
|
else:
|
|
return None
|
|
|
|
def op_gate(self, state):
|
|
if is_non_idle_and_non_empty(
|
|
state.forward_batch.forward_mode, state.hidden_states_mlp_input
|
|
):
|
|
# router_logits: (num_tokens, n_experts)
|
|
state.router_logits = self.gate(state.hidden_states_mlp_input)
|
|
else:
|
|
state.router_logits = None
|
|
|
|
def op_shared_experts(self, state):
|
|
hidden_states_mlp_input = state.pop("hidden_states_mlp_input")
|
|
if (self.num_fused_shared_experts == 0) and is_non_idle_and_non_empty(
|
|
state.forward_batch.forward_mode, hidden_states_mlp_input
|
|
):
|
|
state.shared_output = self.shared_experts(hidden_states_mlp_input)
|
|
else:
|
|
state.shared_output = None
|
|
|
|
def op_select_experts(self, state):
|
|
router_logits = state.pop("router_logits")
|
|
hidden_states = state.hidden_states_mlp_input
|
|
|
|
if router_logits is not None:
|
|
with get_global_expert_distribution_recorder().with_current_layer(
|
|
self.layer_id
|
|
):
|
|
state.topk_output = self.topk(
|
|
hidden_states=hidden_states,
|
|
router_logits=router_logits,
|
|
num_token_non_padded=state.forward_batch.num_token_non_padded,
|
|
expert_location_dispatch_info=ExpertLocationDispatchInfo.init_new(
|
|
layer_id=self.layer_id,
|
|
),
|
|
)
|
|
else:
|
|
state.topk_output = self.topk.empty_topk_output(hidden_states.device)
|
|
|
|
def op_dispatch_a(self, state):
|
|
if self.ep_size > 1:
|
|
self.experts.dispatcher.dispatch_a(
|
|
hidden_states=state.hidden_states_mlp_input,
|
|
topk_output=state.pop("topk_output"),
|
|
tbo_subbatch_index=state.get("tbo_subbatch_index"),
|
|
)
|
|
|
|
def op_dispatch_b(self, state):
|
|
if self.ep_size > 1:
|
|
with get_global_expert_distribution_recorder().with_current_layer(
|
|
self.layer_id
|
|
):
|
|
state.dispatch_output = self.experts.dispatcher.dispatch_b(
|
|
tbo_subbatch_index=state.get("tbo_subbatch_index"),
|
|
)
|
|
|
|
def op_experts(self, state):
|
|
state.combine_input = self.experts.run_moe_core(
|
|
dispatch_output=state.dispatch_output,
|
|
)
|
|
|
|
def op_combine_a(self, state):
|
|
if self.ep_size > 1:
|
|
self.experts.dispatcher.combine_a(
|
|
combine_input=state.pop("combine_input"),
|
|
tbo_subbatch_index=state.get("tbo_subbatch_index"),
|
|
)
|
|
state.pop("dispatch_output")
|
|
|
|
def op_combine_b(self, state):
|
|
if self.ep_size > 1:
|
|
state.hidden_states_after_combine = self.experts.dispatcher.combine_b(
|
|
tbo_subbatch_index=state.get("tbo_subbatch_index"),
|
|
)
|
|
|
|
def op_output(self, state):
|
|
final_hidden_states = state.pop("hidden_states_after_combine")
|
|
|
|
if get_moe_a2a_backend().is_mori():
|
|
num_tokens = state.pop("num_tokens")
|
|
final_hidden_states = final_hidden_states[:num_tokens]
|
|
|
|
if (shared_output := state.pop("shared_output")) is not None:
|
|
x = shared_output
|
|
if _use_aiter:
|
|
x.add_(final_hidden_states)
|
|
else:
|
|
x.add_(final_hidden_states, alpha=self.routed_scaling_factor)
|
|
final_hidden_states = x
|
|
elif _use_aiter:
|
|
# fused in aiter_biased_grouped_topk so we can skip here
|
|
pass
|
|
else:
|
|
final_hidden_states *= self.routed_scaling_factor
|
|
|
|
state.hidden_states_mlp_output = final_hidden_states
|
|
|
|
|
|
class DeepseekV2AttentionMLA(
|
|
nn.Module,
|
|
DeepseekMHAForwardMixin,
|
|
DeepseekMLAForwardMixin,
|
|
DeepseekMLARocmForwardMixin,
|
|
DeepseekMLACpuForwardMixin,
|
|
):
|
|
def __init__(
|
|
self,
|
|
config: PretrainedConfig,
|
|
hidden_size: int,
|
|
num_heads: int,
|
|
qk_nope_head_dim: int,
|
|
qk_rope_head_dim: int,
|
|
v_head_dim: int,
|
|
q_lora_rank: int,
|
|
kv_lora_rank: int,
|
|
rope_theta: float = 10000,
|
|
rope_scaling: Optional[Dict[str, Any]] = None,
|
|
max_position_embeddings: int = 8192,
|
|
quant_config: Optional[QuantizationConfig] = None,
|
|
reduce_results: bool = True,
|
|
layer_id: int = None,
|
|
prefix: str = "",
|
|
alt_stream: Optional[torch.cuda.Stream] = None,
|
|
skip_rope: bool = False,
|
|
) -> None:
|
|
super().__init__()
|
|
self.layer_id = layer_id
|
|
self.hidden_size = hidden_size
|
|
self.qk_nope_head_dim = qk_nope_head_dim
|
|
self.qk_rope_head_dim = qk_rope_head_dim
|
|
self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
|
|
self.v_head_dim = v_head_dim
|
|
self.q_lora_rank = q_lora_rank
|
|
self.kv_lora_rank = kv_lora_rank
|
|
self.quant_config = quant_config
|
|
attn_tp_rank = get_attention_tp_rank()
|
|
attn_tp_size = get_attention_tp_size()
|
|
self.use_nsa = is_deepseek_nsa(config)
|
|
self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp()
|
|
if self.nsa_enable_prefill_cp:
|
|
assert self.use_nsa, "CP currently only supports deepseek v3.2 model"
|
|
# cp reuse the attn_tp comm group but need to duplicate the weights
|
|
if self.nsa_enable_prefill_cp and self.use_nsa:
|
|
self.cp_size = get_attention_cp_size()
|
|
self.num_heads = num_heads
|
|
assert num_heads % attn_tp_size == 0
|
|
self.num_local_heads = num_heads // attn_tp_size
|
|
self.scaling = self.qk_head_dim**-0.5
|
|
self.rope_theta = rope_theta
|
|
self.max_position_embeddings = max_position_embeddings
|
|
self.kv_cache_dtype = get_global_server_args().kv_cache_dtype
|
|
|
|
# NOTE modification to rope_scaling must be done early enough, b/c e.g. Indexer needs it
|
|
if rope_scaling:
|
|
rope_scaling["rope_type"] = "deepseek_yarn"
|
|
|
|
# For tensor parallel attention
|
|
if self.q_lora_rank is not None:
|
|
self.fused_qkv_a_proj_with_mqa = ReplicatedLinear(
|
|
self.hidden_size,
|
|
self.q_lora_rank + self.kv_lora_rank + self.qk_rope_head_dim,
|
|
bias=False,
|
|
quant_config=quant_config,
|
|
prefix=add_prefix("fused_qkv_a_proj_with_mqa", prefix),
|
|
)
|
|
self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps)
|
|
self.q_b_proj = ColumnParallelLinear(
|
|
q_lora_rank,
|
|
self.num_heads * self.qk_head_dim,
|
|
bias=False,
|
|
quant_config=self._get_q_b_proj_quant_config(quant_config),
|
|
prefix=add_prefix("q_b_proj", prefix),
|
|
tp_rank=attn_tp_rank,
|
|
tp_size=attn_tp_size,
|
|
)
|
|
else:
|
|
self.q_proj = ColumnParallelLinear(
|
|
self.hidden_size,
|
|
self.num_heads * self.qk_head_dim,
|
|
bias=False,
|
|
quant_config=quant_config,
|
|
prefix=add_prefix("q_proj", prefix),
|
|
tp_rank=attn_tp_rank,
|
|
tp_size=attn_tp_size,
|
|
)
|
|
self.kv_a_proj_with_mqa = ReplicatedLinear(
|
|
self.hidden_size,
|
|
self.kv_lora_rank + self.qk_rope_head_dim,
|
|
bias=False,
|
|
quant_config=quant_config,
|
|
prefix=add_prefix("kv_a_proj_with_mqa", prefix),
|
|
)
|
|
|
|
if self.use_nsa:
|
|
is_neox_style = not getattr(config, "indexer_rope_interleave", False)
|
|
self.indexer = Indexer(
|
|
hidden_size=hidden_size,
|
|
index_n_heads=get_nsa_index_n_heads(config),
|
|
index_head_dim=get_nsa_index_head_dim(config),
|
|
rope_head_dim=qk_rope_head_dim,
|
|
index_topk=get_nsa_index_topk(config),
|
|
q_lora_rank=q_lora_rank,
|
|
max_position_embeddings=max_position_embeddings,
|
|
rope_theta=rope_theta,
|
|
scale_fmt="ue8m0",
|
|
block_size=128,
|
|
rope_scaling=rope_scaling,
|
|
is_neox_style=is_neox_style,
|
|
prefix=add_prefix("indexer", prefix),
|
|
quant_config=quant_config,
|
|
layer_id=layer_id,
|
|
alt_stream=alt_stream,
|
|
)
|
|
|
|
self.kv_b_proj = ColumnParallelLinear(
|
|
self.kv_lora_rank,
|
|
self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
|
|
bias=False,
|
|
quant_config=quant_config,
|
|
prefix=add_prefix("kv_b_proj", prefix),
|
|
tp_rank=attn_tp_rank,
|
|
tp_size=attn_tp_size,
|
|
)
|
|
# O projection.
|
|
self.o_proj = RowParallelLinear(
|
|
self.num_heads * self.v_head_dim,
|
|
self.hidden_size,
|
|
bias=False,
|
|
quant_config=quant_config,
|
|
reduce_results=reduce_results,
|
|
prefix=add_prefix("o_proj", prefix),
|
|
tp_rank=attn_tp_rank,
|
|
tp_size=attn_tp_size,
|
|
)
|
|
self.kv_a_layernorm = RMSNorm(self.kv_lora_rank, eps=config.rms_norm_eps)
|
|
|
|
if not skip_rope:
|
|
is_neox_style = not getattr(config, "rope_interleave", True)
|
|
self.rotary_emb = get_rope_wrapper(
|
|
qk_rope_head_dim,
|
|
rotary_dim=qk_rope_head_dim,
|
|
max_position=max_position_embeddings,
|
|
base=rope_theta,
|
|
rope_scaling=rope_scaling,
|
|
is_neox_style=is_neox_style,
|
|
device=get_global_server_args().device,
|
|
)
|
|
|
|
if rope_scaling and rope_scaling.get("apply_yarn_scaling", True):
|
|
self.scaling = compute_mla_mscale_scaling(rope_scaling, self.scaling)
|
|
else:
|
|
self.rotary_emb = None
|
|
self.use_deepseek_yarn_rope = rope_scaling is not None
|
|
|
|
self.attn_mqa = RadixAttention(
|
|
self.num_local_heads,
|
|
self.kv_lora_rank + self.qk_rope_head_dim,
|
|
self.scaling,
|
|
num_kv_heads=1,
|
|
layer_id=layer_id,
|
|
v_head_dim=self.kv_lora_rank,
|
|
quant_config=quant_config,
|
|
prefix=add_prefix("attn_mqa", prefix),
|
|
)
|
|
|
|
self.attn_mha = RadixAttention(
|
|
self.num_local_heads,
|
|
self.qk_nope_head_dim + self.qk_rope_head_dim,
|
|
self.scaling,
|
|
num_kv_heads=self.num_local_heads,
|
|
layer_id=layer_id,
|
|
v_head_dim=self.v_head_dim,
|
|
quant_config=quant_config,
|
|
prefix=add_prefix("attn_mha", prefix),
|
|
)
|
|
|
|
self.alt_stream = alt_stream
|
|
self.attn_mha.kv_b_proj = None
|
|
|
|
self.w_kc = None
|
|
self.w_vc = None
|
|
self.w_scale = 1.0
|
|
|
|
self.w_scale_k = None
|
|
self.w_scale_v = None
|
|
self.use_deep_gemm_bmm = False
|
|
|
|
self.current_attention_backend = (
|
|
None # Attention backend used by current forward batch
|
|
)
|
|
|
|
self.has_fused_proj = hasattr(self, "fused_qkv_a_proj_with_mqa")
|
|
self.is_packed_weight = (
|
|
self.has_fused_proj
|
|
and hasattr(self.fused_qkv_a_proj_with_mqa.quant_method, "quant_config")
|
|
and self.fused_qkv_a_proj_with_mqa.quant_method.quant_config.get_name()
|
|
in {"awq", "awq_marlin", "moe_wna16"}
|
|
)
|
|
self.use_min_latency_fused_a_gemm = (
|
|
self.has_fused_proj
|
|
and not self.is_packed_weight
|
|
and self.fused_qkv_a_proj_with_mqa.weight.dtype == torch.bfloat16
|
|
and self.fused_qkv_a_proj_with_mqa.weight.shape[0] == 2112
|
|
and self.fused_qkv_a_proj_with_mqa.weight.shape[1] == 7168
|
|
and _is_cuda
|
|
and 90 <= _device_sm < 120
|
|
)
|
|
|
|
self.init_mha_forward()
|
|
self.init_mla_forward()
|
|
self.init_mla_fused_rope_rocm_forward()
|
|
self.init_mla_fused_rope_cpu_forward()
|
|
|
|
def dispatch_attn_forward_method(
|
|
self, forward_batch: ForwardBatch
|
|
) -> AttnForwardMethod:
|
|
# Determine attention backend used by current forward batch
|
|
if forward_batch.forward_mode.is_decode_or_idle():
|
|
attention_backend = get_global_server_args().decode_attention_backend
|
|
elif (
|
|
forward_batch.forward_mode.is_target_verify()
|
|
or forward_batch.forward_mode.is_draft_extend()
|
|
):
|
|
# Use the specified backend for speculative operations (both verify and draft extend)
|
|
if get_global_server_args().speculative_attention_mode == "decode":
|
|
attention_backend = get_global_server_args().decode_attention_backend
|
|
else: # default to prefill
|
|
attention_backend = get_global_server_args().prefill_attention_backend
|
|
else:
|
|
attention_backend = get_global_server_args().prefill_attention_backend
|
|
self.current_attention_backend = attention_backend
|
|
|
|
handler = AttentionBackendRegistry.get_handler(attention_backend)
|
|
return handler(self, forward_batch)
|
|
|
|
def op_prepare(self, state):
|
|
state.attn_intermediate_state = self.forward_prepare(
|
|
positions=state.positions,
|
|
hidden_states=state.pop("hidden_states_after_comm_pre_attn"),
|
|
forward_batch=state.forward_batch,
|
|
zero_allocator=state.zero_allocator,
|
|
)
|
|
|
|
def op_core(self, state):
|
|
state.hidden_states_after_attn = self.forward_core(
|
|
state.pop("attn_intermediate_state")
|
|
)
|
|
|
|
def forward(
|
|
self,
|
|
positions: torch.Tensor,
|
|
hidden_states: torch.Tensor,
|
|
forward_batch: ForwardBatch,
|
|
zero_allocator: BumpAllocator,
|
|
layer_scatter_modes: LayerScatterModes = None,
|
|
llama_4_scaling: Optional[torch.Tensor] = None,
|
|
):
|
|
s = self.forward_prepare(
|
|
positions=positions,
|
|
hidden_states=hidden_states,
|
|
forward_batch=forward_batch,
|
|
zero_allocator=zero_allocator,
|
|
layer_scatter_modes=layer_scatter_modes,
|
|
llama_4_scaling=llama_4_scaling,
|
|
)
|
|
return self.forward_core(s)
|
|
|
|
def forward_prepare(
|
|
self,
|
|
positions: torch.Tensor,
|
|
hidden_states: torch.Tensor,
|
|
forward_batch: ForwardBatch,
|
|
zero_allocator: BumpAllocator,
|
|
layer_scatter_modes: LayerScatterModes = None,
|
|
llama_4_scaling: Optional[torch.Tensor] = None,
|
|
):
|
|
if self.attn_mha.kv_b_proj is None:
|
|
self.attn_mha.kv_b_proj = self.kv_b_proj
|
|
|
|
# when hidden_states is a tuple of tensors, the tuple will include quantized weight and scale tensor
|
|
if isinstance(hidden_states, tuple):
|
|
if (
|
|
not get_attn_tp_context().input_scattered
|
|
and hidden_states[0].shape[0] == 0
|
|
):
|
|
assert not self.o_proj.reduce_results, (
|
|
"short-circuiting allreduce will lead to hangs"
|
|
)
|
|
return hidden_states[0]
|
|
else:
|
|
if (
|
|
not get_attn_tp_context().input_scattered
|
|
and hidden_states.shape[0] == 0
|
|
):
|
|
assert not self.o_proj.reduce_results, (
|
|
"short-circuiting allreduce will lead to hangs"
|
|
)
|
|
return hidden_states, None, forward_batch, None
|
|
|
|
attn_forward_method = self.dispatch_attn_forward_method(forward_batch)
|
|
if attn_forward_method == AttnForwardMethod.MHA:
|
|
inner_state = self.forward_normal_prepare(
|
|
positions, hidden_states, forward_batch, zero_allocator
|
|
)
|
|
elif attn_forward_method == AttnForwardMethod.MHA_CHUNKED_KV:
|
|
inner_state = self.forward_normal_chunked_kv_prepare(
|
|
positions, hidden_states, forward_batch, zero_allocator
|
|
)
|
|
elif attn_forward_method == AttnForwardMethod.MHA_ONE_SHOT:
|
|
inner_state = self.forward_normal_one_shot_prepare(
|
|
positions, hidden_states, forward_batch, zero_allocator
|
|
)
|
|
elif attn_forward_method == AttnForwardMethod.MLA:
|
|
inner_state = self.forward_absorb_prepare(
|
|
positions, hidden_states, forward_batch, zero_allocator, llama_4_scaling
|
|
)
|
|
elif attn_forward_method == AttnForwardMethod.MLA_FUSED_ROPE_ROCM:
|
|
inner_state = self.forward_absorb_fused_mla_rope_prepare(
|
|
positions, hidden_states, forward_batch, zero_allocator
|
|
)
|
|
elif attn_forward_method == AttnForwardMethod.MLA_FUSED_ROPE_CPU:
|
|
inner_state = self.forward_absorb_fused_mla_rope_cpu_prepare(
|
|
positions, hidden_states, forward_batch, zero_allocator
|
|
)
|
|
elif attn_forward_method == AttnForwardMethod.MHA_NPU:
|
|
inner_state = forward_mha_prepare_npu(
|
|
self,
|
|
positions,
|
|
hidden_states,
|
|
forward_batch,
|
|
zero_allocator,
|
|
layer_scatter_modes,
|
|
)
|
|
elif attn_forward_method == AttnForwardMethod.MLA_NPU:
|
|
inner_state = forward_mla_prepare_npu(
|
|
self,
|
|
positions,
|
|
hidden_states,
|
|
forward_batch,
|
|
zero_allocator,
|
|
layer_scatter_modes,
|
|
)
|
|
elif attn_forward_method == AttnForwardMethod.DSA_NPU:
|
|
inner_state = forward_dsa_prepare_npu(
|
|
self,
|
|
positions,
|
|
hidden_states,
|
|
forward_batch,
|
|
zero_allocator,
|
|
layer_scatter_modes,
|
|
)
|
|
else:
|
|
raise NotImplementedError
|
|
return None, attn_forward_method, forward_batch, inner_state
|
|
|
|
def forward_core(self, intermediate_state):
|
|
hidden_states, attn_forward_method, forward_batch, inner_state = (
|
|
intermediate_state
|
|
)
|
|
if inner_state is None:
|
|
return hidden_states
|
|
|
|
if attn_forward_method == AttnForwardMethod.MHA:
|
|
return self.forward_normal_core(*inner_state)
|
|
elif attn_forward_method == AttnForwardMethod.MHA_CHUNKED_KV:
|
|
return self.forward_normal_chunked_kv_core(*inner_state)
|
|
elif attn_forward_method == AttnForwardMethod.MHA_ONE_SHOT:
|
|
return self.forward_normal_one_shot_core(*inner_state)
|
|
elif attn_forward_method == AttnForwardMethod.MLA:
|
|
return self.forward_absorb_core(*inner_state)
|
|
elif attn_forward_method == AttnForwardMethod.MLA_FUSED_ROPE_ROCM:
|
|
return self.forward_absorb_fused_mla_rope_core(*inner_state)
|
|
elif attn_forward_method == AttnForwardMethod.MLA_FUSED_ROPE_CPU:
|
|
return self.forward_absorb_fused_mla_rope_cpu_core(*inner_state)
|
|
elif attn_forward_method == AttnForwardMethod.MHA_NPU:
|
|
return forward_mha_core_npu(self, *inner_state)
|
|
elif attn_forward_method == AttnForwardMethod.MLA_NPU:
|
|
return forward_mla_core_npu(self, *inner_state)
|
|
elif attn_forward_method == AttnForwardMethod.DSA_NPU:
|
|
return forward_dsa_core_npu(self, *inner_state)
|
|
else:
|
|
raise NotImplementedError
|
|
|
|
def prepare_qkv_latent(
|
|
self, hidden_states: torch.Tensor, forward_batch: ForwardBatch
|
|
):
|
|
assert self.q_lora_rank is not None
|
|
if (
|
|
(not isinstance(hidden_states, tuple))
|
|
and hidden_states.shape[0] >= 1
|
|
and hidden_states.shape[0] <= 16
|
|
and self.use_min_latency_fused_a_gemm
|
|
):
|
|
qkv_latent = dsv3_fused_a_gemm(
|
|
hidden_states, self.fused_qkv_a_proj_with_mqa.weight.T
|
|
)
|
|
else:
|
|
qkv_latent = self.fused_qkv_a_proj_with_mqa(hidden_states)[0]
|
|
return qkv_latent
|
|
|
|
def rebuild_cp_kv_cache(self, latent_cache, forward_batch, k_nope, k_pe):
|
|
# support allgather+rerrange
|
|
latent_cache[..., : self.kv_lora_rank] = k_nope.squeeze(1)
|
|
latent_cache[..., self.kv_lora_rank :] = k_pe.squeeze(1)
|
|
latent_cache_output = cp_all_gather_rerange_output(
|
|
latent_cache.contiguous(),
|
|
self.cp_size,
|
|
forward_batch,
|
|
torch.cuda.current_stream(),
|
|
)
|
|
k_nope = latent_cache_output[..., : self.kv_lora_rank].unsqueeze(1)
|
|
k_pe = latent_cache_output[..., self.kv_lora_rank :].unsqueeze(1)
|
|
return k_nope, k_pe
|
|
|
|
@staticmethod
|
|
def _get_q_b_proj_quant_config(quant_config):
|
|
if envs.SGLANG_NVFP4_CKPT_FP8_GEMM_IN_ATTN.get():
|
|
# refer to real DeepSeek V3 quant config
|
|
return Fp8Config(
|
|
is_checkpoint_fp8_serialized=True,
|
|
weight_block_size=[128, 128],
|
|
)
|
|
else:
|
|
return quant_config
|
|
|
|
|
|
class DeepseekV2DecoderLayer(nn.Module):
|
|
def __init__(
|
|
self,
|
|
config: PretrainedConfig,
|
|
layer_id: int,
|
|
quant_config: Optional[QuantizationConfig] = None,
|
|
moe_quant_config_override: Optional[QuantizationConfig] = None,
|
|
is_nextn: bool = False,
|
|
prefix: str = "",
|
|
alt_stream: Optional[torch.cuda.Stream] = None,
|
|
) -> None:
|
|
super().__init__()
|
|
self.hidden_size = config.hidden_size
|
|
self.config = config
|
|
if hasattr(config, "rope_parameters"):
|
|
rope_theta = config.rope_parameters["rope_theta"]
|
|
assert rope_theta is not None, f"rope_theta not found in config: {config}"
|
|
rope_type = config.rope_parameters.get("rope_type")
|
|
rope_scaling = config.rope_parameters if rope_type != "default" else None
|
|
else:
|
|
rope_theta = config.rope_theta
|
|
rope_scaling = config.rope_scaling
|
|
max_position_embeddings = config.max_position_embeddings
|
|
self.speculative_algorithm = SpeculativeAlgorithm.from_string(
|
|
get_global_server_args().speculative_algorithm
|
|
)
|
|
self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp()
|
|
self.layer_id = layer_id
|
|
self.is_nextn = is_nextn
|
|
self.self_attn = DeepseekV2AttentionMLA(
|
|
config=config,
|
|
hidden_size=self.hidden_size,
|
|
num_heads=config.num_attention_heads,
|
|
qk_nope_head_dim=config.qk_nope_head_dim,
|
|
qk_rope_head_dim=config.qk_rope_head_dim,
|
|
v_head_dim=config.v_head_dim,
|
|
q_lora_rank=(
|
|
config.q_lora_rank if hasattr(config, "q_lora_rank") else None
|
|
),
|
|
kv_lora_rank=config.kv_lora_rank,
|
|
rope_theta=rope_theta,
|
|
rope_scaling=rope_scaling,
|
|
max_position_embeddings=max_position_embeddings,
|
|
quant_config=quant_config,
|
|
layer_id=layer_id,
|
|
reduce_results=False,
|
|
prefix=add_prefix("self_attn", prefix),
|
|
alt_stream=alt_stream,
|
|
)
|
|
if not hasattr(config, "q_lora_rank") and envs.SGLANG_USE_AG_AFTER_QLORA.get():
|
|
raise ValueError(
|
|
"SGLANG_USE_AG_AFTER_QLORA only supports the model with q_lora_rank"
|
|
)
|
|
|
|
self.is_layer_sparse = self._is_layer_sparse(layer_id, is_nextn=is_nextn)
|
|
is_previous_layer_sparse = self._is_layer_sparse(layer_id - 1, is_nextn=False)
|
|
is_next_layer_sparse = self._is_layer_sparse(layer_id + 1, is_nextn=False)
|
|
|
|
self.layer_scatter_modes = LayerScatterModes.init_new(
|
|
layer_id=layer_id,
|
|
num_layers=1 if is_nextn else config.num_hidden_layers,
|
|
is_layer_sparse=self.is_layer_sparse,
|
|
is_previous_layer_sparse=is_previous_layer_sparse,
|
|
is_next_layer_sparse=is_next_layer_sparse,
|
|
)
|
|
|
|
if self.is_layer_sparse:
|
|
self.mlp = DeepseekV2MoE(
|
|
config=config,
|
|
quant_config=moe_quant_config_override or quant_config,
|
|
prefix=add_prefix("mlp", prefix),
|
|
layer_id=self.layer_id,
|
|
alt_stream=alt_stream,
|
|
is_nextn=is_nextn,
|
|
)
|
|
else:
|
|
if enable_moe_dense_fully_dp():
|
|
mlp_tp_rank, mlp_tp_size = 0, 1
|
|
else:
|
|
mlp_tp_rank, mlp_tp_size = None, None
|
|
self.mlp = DeepseekV2MLP(
|
|
hidden_size=config.hidden_size,
|
|
intermediate_size=config.intermediate_size,
|
|
hidden_act=config.hidden_act,
|
|
quant_config=quant_config,
|
|
prefix=add_prefix("mlp", prefix),
|
|
tp_rank=mlp_tp_rank,
|
|
tp_size=mlp_tp_size,
|
|
)
|
|
|
|
self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
|
self.post_attention_layernorm = RMSNorm(
|
|
config.hidden_size, eps=config.rms_norm_eps
|
|
)
|
|
|
|
if self.nsa_enable_prefill_cp:
|
|
self.layer_communicator = NSACPLayerCommunicator(
|
|
layer_scatter_modes=self.layer_scatter_modes,
|
|
input_layernorm=self.input_layernorm,
|
|
post_attention_layernorm=self.post_attention_layernorm,
|
|
allow_reduce_scatter=True,
|
|
is_last_layer=(
|
|
is_nextn or (self.layer_id == self.config.num_hidden_layers - 1)
|
|
),
|
|
qkv_latent_func=self.self_attn.prepare_qkv_latent,
|
|
)
|
|
else:
|
|
self.layer_communicator = LayerCommunicator(
|
|
layer_scatter_modes=self.layer_scatter_modes,
|
|
input_layernorm=self.input_layernorm,
|
|
post_attention_layernorm=self.post_attention_layernorm,
|
|
allow_reduce_scatter=True,
|
|
is_last_layer=(
|
|
is_nextn or (self.layer_id == self.config.num_hidden_layers - 1)
|
|
),
|
|
qkv_latent_func=self.self_attn.prepare_qkv_latent,
|
|
)
|
|
|
|
def _is_layer_sparse(self, layer_id: int, is_nextn: bool) -> bool:
|
|
return is_nextn or (
|
|
self.config.n_routed_experts is not None
|
|
and layer_id >= self.config.first_k_dense_replace
|
|
and layer_id % self.config.moe_layer_freq == 0
|
|
)
|
|
|
|
def _notify_cp_hicache_layer_end(
|
|
self,
|
|
forward_batch: ForwardBatch,
|
|
tbo_subbatch_index: Optional[int] = None,
|
|
) -> None:
|
|
if (
|
|
getattr(forward_batch, "tbo_parent_token_range", None) is not None
|
|
and tbo_subbatch_index != 1
|
|
):
|
|
return
|
|
notifier = getattr(
|
|
forward_batch.token_to_kv_pool, "notify_layer_end_for_backup", None
|
|
)
|
|
if notifier is not None:
|
|
notifier(self.layer_id)
|
|
|
|
def forward(
|
|
self,
|
|
positions: torch.Tensor,
|
|
hidden_states: torch.Tensor,
|
|
forward_batch: ForwardBatch,
|
|
residual: Optional[torch.Tensor],
|
|
zero_allocator: BumpAllocator,
|
|
gemm_output_zero_allocator: BumpAllocator = None,
|
|
llama_4_scaling: Optional[torch.Tensor] = None,
|
|
) -> torch.Tensor:
|
|
quant_format = (
|
|
"mxfp4"
|
|
if (
|
|
_is_gfx95_supported
|
|
and getattr(self.self_attn, "fused_qkv_a_proj_with_mqa", None)
|
|
is not None
|
|
and getattr(self.self_attn.fused_qkv_a_proj_with_mqa, "weight", None)
|
|
is not None
|
|
and self.self_attn.fused_qkv_a_proj_with_mqa.weight.dtype == torch.uint8
|
|
)
|
|
else (
|
|
"fp8"
|
|
if (
|
|
_is_gfx95_supported
|
|
and getattr(self.self_attn, "fused_qkv_a_proj_with_mqa", None)
|
|
is not None
|
|
and getattr(
|
|
self.self_attn.fused_qkv_a_proj_with_mqa, "weight", None
|
|
)
|
|
is not None
|
|
and self.self_attn.fused_qkv_a_proj_with_mqa.weight.dtype
|
|
== getattr(torch, "float8_e4m3fn", None)
|
|
)
|
|
else ""
|
|
)
|
|
)
|
|
|
|
hidden_states, residual = self.layer_communicator.prepare_attn(
|
|
hidden_states,
|
|
residual,
|
|
forward_batch,
|
|
quant_format,
|
|
)
|
|
|
|
hidden_states = self.self_attn(
|
|
positions=positions,
|
|
hidden_states=hidden_states,
|
|
forward_batch=forward_batch,
|
|
zero_allocator=zero_allocator,
|
|
llama_4_scaling=llama_4_scaling,
|
|
layer_scatter_modes=self.layer_scatter_modes,
|
|
)
|
|
|
|
hidden_states, residual = self.layer_communicator.prepare_mlp(
|
|
hidden_states, residual, forward_batch
|
|
)
|
|
|
|
should_allreduce_fusion = (
|
|
self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer(
|
|
forward_batch
|
|
)
|
|
)
|
|
|
|
# For DP with padding, reduce scatter can be used instead of all-reduce.
|
|
use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter(
|
|
forward_batch
|
|
)
|
|
|
|
if isinstance(self.mlp, DeepseekV2MLP):
|
|
gemm_output_zero_allocator = None
|
|
|
|
hidden_states = self.mlp(
|
|
hidden_states,
|
|
forward_batch,
|
|
should_allreduce_fusion,
|
|
use_reduce_scatter,
|
|
gemm_output_zero_allocator,
|
|
)
|
|
|
|
if not self.nsa_enable_prefill_cp and should_allreduce_fusion:
|
|
hidden_states._sglang_needs_allreduce_fusion = True
|
|
|
|
if not should_allreduce_fusion:
|
|
hidden_states, residual = self.layer_communicator.postprocess_layer(
|
|
hidden_states, residual, forward_batch
|
|
)
|
|
|
|
self._notify_cp_hicache_layer_end(forward_batch)
|
|
return hidden_states, residual
|
|
|
|
def op_comm_prepare_attn(
|
|
self,
|
|
state,
|
|
positions: torch.Tensor,
|
|
hidden_states: torch.Tensor,
|
|
forward_batch: ForwardBatch,
|
|
residual: Optional[torch.Tensor],
|
|
zero_allocator: BumpAllocator,
|
|
tbo_subbatch_index: Optional[int] = None,
|
|
):
|
|
state.hidden_states_after_comm_pre_attn, state.residual_after_input_ln = (
|
|
self.layer_communicator.prepare_attn(hidden_states, residual, forward_batch)
|
|
)
|
|
if get_moe_a2a_backend().is_mori():
|
|
state.num_tokens = hidden_states.shape[0]
|
|
state.update(
|
|
dict(
|
|
forward_batch=forward_batch,
|
|
positions=positions,
|
|
zero_allocator=zero_allocator,
|
|
tbo_subbatch_index=tbo_subbatch_index,
|
|
)
|
|
)
|
|
|
|
def op_comm_prepare_mlp(self, state):
|
|
state.hidden_states_mlp_input, state.residual_after_comm_pre_mlp = (
|
|
self.layer_communicator.prepare_mlp(
|
|
state.pop("hidden_states_after_attn"),
|
|
state.pop("residual_after_input_ln"),
|
|
state.forward_batch,
|
|
)
|
|
)
|
|
|
|
def op_mlp(self, state):
|
|
hidden_states = state.pop("hidden_states_mlp_input")
|
|
if not (
|
|
enable_moe_dense_fully_dp()
|
|
and (not self.is_layer_sparse)
|
|
and hidden_states.shape[0] == 0
|
|
):
|
|
state.hidden_states_mlp_output = self.mlp(
|
|
hidden_states, state.forward_batch
|
|
)
|
|
else:
|
|
state.hidden_states_mlp_output = hidden_states
|
|
|
|
def op_comm_postprocess_layer(self, state):
|
|
hidden_states, residual = self.layer_communicator.postprocess_layer(
|
|
state.pop("hidden_states_mlp_output"),
|
|
state.pop("residual_after_comm_pre_mlp"),
|
|
state.forward_batch,
|
|
)
|
|
self._notify_cp_hicache_layer_end(
|
|
state.forward_batch, tbo_subbatch_index=state.tbo_subbatch_index
|
|
)
|
|
|
|
output = dict(
|
|
positions=state.positions,
|
|
hidden_states=hidden_states,
|
|
residual=residual,
|
|
forward_batch=state.forward_batch,
|
|
zero_allocator=state.zero_allocator,
|
|
tbo_subbatch_index=state.tbo_subbatch_index,
|
|
)
|
|
|
|
state.clear(
|
|
expect_keys={
|
|
"positions",
|
|
"forward_batch",
|
|
"zero_allocator",
|
|
"tbo_subbatch_index",
|
|
}
|
|
)
|
|
return output
|
|
|
|
|
|
class DeepseekV2Model(nn.Module):
|
|
fall_back_to_pt_during_load = False
|
|
|
|
def __init__(
|
|
self,
|
|
config: PretrainedConfig,
|
|
quant_config: Optional[QuantizationConfig] = None,
|
|
prefix: str = "",
|
|
) -> None:
|
|
super().__init__()
|
|
self.padding_id = config.pad_token_id
|
|
self.vocab_size = config.vocab_size
|
|
self.first_k_dense_replace = config.first_k_dense_replace
|
|
self.pp_group = get_pp_group()
|
|
self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp()
|
|
if self.nsa_enable_prefill_cp:
|
|
self.cp_size = get_attention_cp_size()
|
|
else:
|
|
self.cp_size = None
|
|
|
|
if self.pp_group.is_first_rank:
|
|
self.embed_tokens = VocabParallelEmbedding(
|
|
config.vocab_size,
|
|
config.hidden_size,
|
|
use_attn_tp_group=is_dp_attention_enabled(),
|
|
)
|
|
else:
|
|
self.embed_tokens = PPMissingLayer()
|
|
|
|
self.alt_stream = (
|
|
torch.cuda.Stream()
|
|
if _is_cuda or envs.SGLANG_NPU_USE_MULTI_STREAM.get()
|
|
else None
|
|
)
|
|
|
|
self.layers, self.start_layer, self.end_layer = make_layers(
|
|
config.num_hidden_layers,
|
|
lambda idx, prefix: DeepseekV2DecoderLayer(
|
|
config=config,
|
|
layer_id=idx,
|
|
quant_config=quant_config,
|
|
prefix=prefix,
|
|
alt_stream=self.alt_stream,
|
|
),
|
|
pp_rank=self.pp_group.rank_in_group,
|
|
pp_size=self.pp_group.world_size,
|
|
prefix=add_prefix("layers", prefix),
|
|
offloader_kwargs=dict(
|
|
submodule_accessor=lambda layer: (
|
|
layer.mlp.experts
|
|
if isinstance(layer.mlp, DeepseekV2MoE)
|
|
else layer.mlp
|
|
),
|
|
whitelist_param_names_creator=lambda module: (
|
|
[
|
|
"w13_weight",
|
|
"w2_weight",
|
|
# only for nvfp4
|
|
*(
|
|
[
|
|
"w13_blockscale_swizzled",
|
|
"w2_blockscale_swizzled",
|
|
]
|
|
if hasattr(module, "w13_blockscale_swizzled")
|
|
else []
|
|
),
|
|
]
|
|
if isinstance(module, FusedMoE)
|
|
else []
|
|
),
|
|
),
|
|
)
|
|
if self.pp_group.is_last_rank:
|
|
self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
|
else:
|
|
self.norm = PPMissingLayer(return_tuple=True)
|
|
|
|
self.gemm_output_zero_allocator_size = 0
|
|
if (
|
|
_use_aiter_gfx95
|
|
and config.n_routed_experts == 256
|
|
and self.embed_tokens.embedding_dim == 7168
|
|
):
|
|
num_moe_layers = sum(
|
|
[
|
|
1
|
|
for i in range(len(self.layers))
|
|
if isinstance(self.layers[i].mlp, DeepseekV2MoE)
|
|
]
|
|
)
|
|
|
|
allocate_size = 0
|
|
for i in range(len(self.layers)):
|
|
if isinstance(self.layers[i].mlp, DeepseekV2MoE):
|
|
# tp_size = get_tensor_model_parallel_world_size()
|
|
a2a_backend = get_moe_a2a_backend()
|
|
is_a2a_moe = (
|
|
a2a_backend.is_deepep()
|
|
or a2a_backend.is_mori()
|
|
or a2a_backend.is_mooncake()
|
|
)
|
|
tp_size = (
|
|
1 if is_a2a_moe else get_tensor_model_parallel_world_size()
|
|
)
|
|
intermediate_size = (
|
|
config.moe_intermediate_size * config.n_shared_experts
|
|
)
|
|
share_expert_output_size_per_partition = divide(
|
|
intermediate_size * 2, tp_size
|
|
)
|
|
allocate_size = share_expert_output_size_per_partition
|
|
break
|
|
|
|
self.gemm_output_zero_allocator_size = (
|
|
get_dsv3_gemm_output_zero_allocator_size(
|
|
config.n_routed_experts,
|
|
num_moe_layers,
|
|
allocate_size,
|
|
self.embed_tokens.embedding_dim,
|
|
)
|
|
)
|
|
self.layers_to_capture = []
|
|
if get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mooncake():
|
|
self.enable_a2a_moe = True
|
|
else:
|
|
self.enable_a2a_moe = False
|
|
|
|
# llama_4_scaling: for supporting Mistral-Large-3 model
|
|
self.llama_4_scaling_config = getattr(config, "llama_4_scaling", None)
|
|
|
|
def get_input_embeddings(self) -> torch.Tensor:
|
|
return self.embed_tokens
|
|
|
|
def forward(
|
|
self,
|
|
input_ids: torch.Tensor,
|
|
positions: torch.Tensor,
|
|
forward_batch: ForwardBatch,
|
|
input_embeds: torch.Tensor = None,
|
|
pp_proxy_tensors: Optional[PPProxyTensors] = None,
|
|
) -> Union[torch.Tensor, PPProxyTensors]:
|
|
total_num_layers = self.end_layer - self.start_layer
|
|
if self.pp_group.is_first_rank:
|
|
if input_embeds is None:
|
|
hidden_states = self.embed_tokens(input_ids)
|
|
else:
|
|
hidden_states = input_embeds
|
|
residual = None
|
|
else:
|
|
assert pp_proxy_tensors is not None
|
|
hidden_states = pp_proxy_tensors["hidden_states"]
|
|
residual = pp_proxy_tensors["residual"]
|
|
device = hidden_states.device
|
|
zero_allocator = BumpAllocator(
|
|
buffer_size=total_num_layers * 2 * (2 if forward_batch.can_run_tbo else 1),
|
|
dtype=torch.float32,
|
|
device=device,
|
|
)
|
|
|
|
has_gemm_output_zero_allocator = hasattr(
|
|
self, "gemm_output_zero_allocator_size"
|
|
)
|
|
|
|
gemm_output_zero_allocator = (
|
|
BumpAllocator(
|
|
buffer_size=self.gemm_output_zero_allocator_size,
|
|
dtype=torch.float32,
|
|
device=device,
|
|
)
|
|
if has_gemm_output_zero_allocator
|
|
and self.gemm_output_zero_allocator_size > 0
|
|
else None
|
|
)
|
|
|
|
if nsa_use_prefill_cp(forward_batch):
|
|
if self.pp_group.is_first_rank:
|
|
hidden_states = cp_split_and_rebuild_data(forward_batch, hidden_states)
|
|
positions = cp_split_and_rebuild_position(forward_batch, positions)
|
|
|
|
# llama_4_scaling: for supporting Mistral-Large-3 model
|
|
# Compute llama 4 scaling once per forward pass if enabled
|
|
llama_4_scaling: Optional[torch.Tensor] = None
|
|
if self.llama_4_scaling_config is not None:
|
|
llama_4_scaling = _get_llama_4_scaling(
|
|
original_max_position_embeddings=self.llama_4_scaling_config[
|
|
"original_max_position_embeddings"
|
|
],
|
|
scaling_beta=self.llama_4_scaling_config["beta"],
|
|
positions=positions,
|
|
)
|
|
|
|
normal_start_layer = self.start_layer
|
|
normal_end_layer = self.end_layer
|
|
if forward_batch.can_run_tbo:
|
|
if (
|
|
self.first_k_dense_replace > normal_start_layer
|
|
and self.first_k_dense_replace < normal_end_layer
|
|
):
|
|
normal_end_layer = self.first_k_dense_replace
|
|
elif self.first_k_dense_replace < normal_start_layer:
|
|
normal_end_layer = normal_start_layer = 0
|
|
aux_hidden_states = []
|
|
for i in range(normal_start_layer, normal_end_layer):
|
|
# NOTE: torch dynamo does not support graph break in context manager
|
|
ctx = (
|
|
nullcontext()
|
|
if not get_global_server_args().disable_piecewise_cuda_graph
|
|
else get_global_expert_distribution_recorder().with_current_layer(i)
|
|
)
|
|
with ctx:
|
|
if i in self.layers_to_capture:
|
|
if self.enable_a2a_moe and i > self.first_k_dense_replace:
|
|
aux_hidden_state = get_attention_tp_group().all_gather(
|
|
hidden_states + residual, dim=0
|
|
)
|
|
aux_hidden_states.append(aux_hidden_state)
|
|
else:
|
|
aux_hidden_states.append(hidden_states + residual)
|
|
layer = self.layers[i]
|
|
hidden_states, residual = layer(
|
|
positions,
|
|
hidden_states,
|
|
forward_batch,
|
|
residual,
|
|
zero_allocator,
|
|
gemm_output_zero_allocator,
|
|
llama_4_scaling,
|
|
)
|
|
|
|
if normal_end_layer != self.end_layer:
|
|
hidden_states, residual = model_forward_maybe_tbo(
|
|
layers=self.layers[normal_end_layer : self.end_layer],
|
|
enable_tbo=True,
|
|
positions=positions,
|
|
forward_batch=forward_batch,
|
|
hidden_states=hidden_states,
|
|
residual=residual,
|
|
input_data_scatter_mode=self.layers[
|
|
normal_end_layer - 1
|
|
].layer_scatter_modes.layer_output_mode,
|
|
zero_allocator=zero_allocator,
|
|
)
|
|
|
|
if not self.pp_group.is_last_rank:
|
|
return PPProxyTensors(
|
|
{
|
|
"hidden_states": hidden_states,
|
|
"residual": residual,
|
|
}
|
|
)
|
|
else:
|
|
if not forward_batch.forward_mode.is_idle():
|
|
if residual is None:
|
|
hidden_states = self.norm(hidden_states)
|
|
else:
|
|
hidden_states, _ = self.norm(hidden_states, residual)
|
|
|
|
if getattr(forward_batch, "capture_draft_hidden_states", False):
|
|
forward_batch.draft_hidden_states = hidden_states
|
|
|
|
if self.pp_group.is_last_rank and nsa_use_prefill_cp(forward_batch):
|
|
if self._should_use_narrow_output_path(forward_batch):
|
|
hidden_states = cp_collect_last_token_hidden(
|
|
hidden_states, forward_batch, self.cp_size
|
|
)
|
|
else:
|
|
hidden_states = cp_all_gather_rerange_output(
|
|
hidden_states,
|
|
self.cp_size,
|
|
forward_batch,
|
|
torch.cuda.current_stream(),
|
|
)
|
|
if len(aux_hidden_states) == 0:
|
|
return hidden_states
|
|
return hidden_states, aux_hidden_states
|
|
|
|
def _should_use_narrow_output_path(self, forward_batch):
|
|
if not nsa_use_prefill_cp(forward_batch):
|
|
return False
|
|
if not self.pp_group.is_last_rank:
|
|
return False
|
|
if not forward_batch.forward_mode.is_extend():
|
|
return False
|
|
if forward_batch.return_logprob:
|
|
return False
|
|
if forward_batch.capture_hidden_mode != CaptureHiddenMode.NULL:
|
|
return False
|
|
return True
|
|
|
|
|
|
class DeepseekV2ForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin):
|
|
# for quark model load
|
|
packed_modules_mapping = {}
|
|
|
|
def __init__(
|
|
self,
|
|
config: PretrainedConfig,
|
|
quant_config: Optional[QuantizationConfig] = None,
|
|
prefix: str = "",
|
|
) -> None:
|
|
super().__init__()
|
|
|
|
# for quark model load
|
|
# Fuse q_a_proj and kv_a_proj_with_mqa along output dimension when q_lora_rank is not None
|
|
self.fuse_qkv_a_proj = (
|
|
hasattr(config, "q_lora_rank") and config.q_lora_rank is not None
|
|
)
|
|
if self.fuse_qkv_a_proj:
|
|
self.packed_modules_mapping["fused_qkv_a_proj_with_mqa"] = [
|
|
"q_a_proj",
|
|
"kv_a_proj_with_mqa",
|
|
]
|
|
|
|
self.pp_group = get_pp_group()
|
|
self.config = config
|
|
self.tp_size = get_tensor_model_parallel_world_size()
|
|
self.quant_config = quant_config
|
|
self.determine_num_fused_shared_experts()
|
|
self.use_nsa = is_deepseek_nsa(config)
|
|
self.model = DeepseekV2Model(
|
|
config, quant_config, prefix=add_prefix("model", prefix)
|
|
)
|
|
|
|
if self.pp_group.is_last_rank:
|
|
if self.pp_group.world_size == 1 and config.tie_word_embeddings:
|
|
self.lm_head = self.model.embed_tokens
|
|
else:
|
|
self.lm_head = ParallelLMHead(
|
|
config.vocab_size,
|
|
config.hidden_size,
|
|
quant_config=quant_config,
|
|
prefix=add_prefix("lm_head", prefix),
|
|
use_attn_tp_group=get_global_server_args().enable_dp_lm_head,
|
|
)
|
|
else:
|
|
# ranks other than the last rank will have a placeholder layer
|
|
self.lm_head = PPMissingLayer()
|
|
self.logits_processor = LogitsProcessor(config)
|
|
|
|
self._routed_experts_weights_of_layer = LazyValue(
|
|
lambda: {
|
|
layer_id: layer.mlp.get_moe_weights()
|
|
for layer_id, layer in enumerate(self.model.layers)
|
|
if isinstance(layer.mlp, DeepseekV2MoE)
|
|
}
|
|
)
|
|
self.capture_aux_hidden_states = False
|
|
|
|
self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp()
|
|
if self.nsa_enable_prefill_cp:
|
|
self.cp_rank = get_attention_cp_rank()
|
|
self.cp_size = get_attention_cp_size()
|
|
else:
|
|
self.cp_rank = self.cp_size = None
|
|
|
|
q_lora_rank = config.q_lora_rank if hasattr(config, "q_lora_rank") else None
|
|
get_attn_tp_context().init_context(q_lora_rank, is_deepseek_nsa(config))
|
|
|
|
@property
|
|
def routed_experts_weights_of_layer(self):
|
|
return self._routed_experts_weights_of_layer.value
|
|
|
|
def determine_num_fused_shared_experts(
|
|
self, architecture: str = "DeepseekV3ForCausalLM"
|
|
):
|
|
self.num_fused_shared_experts = 0
|
|
if get_global_server_args().disable_shared_experts_fusion:
|
|
return
|
|
|
|
# Only Deepseek V3/R1 can use shared experts fusion optimization now.
|
|
disable_reason = None
|
|
if (
|
|
self.config.architectures[0] != architecture
|
|
or self.config.n_routed_experts != 256
|
|
or self.config.n_shared_experts != 1
|
|
):
|
|
disable_reason = "Config does not support fused shared expert(s)."
|
|
elif (not _is_cuda or torch.cuda.get_device_capability("cuda") < (8, 0)) and (
|
|
not _is_hip or torch.cuda.get_device_capability("cuda") < (9, 4)
|
|
):
|
|
disable_reason = (
|
|
"Only Deepseek V3/R1 on NV-platform with capability >= 80 "
|
|
"or AMD-platform with capability >= gfx942(MI30x) can use shared experts fusion optimization."
|
|
)
|
|
elif get_moe_expert_parallel_world_size() > 1 and (
|
|
not _is_hip or torch.cuda.get_device_capability("cuda") < (9, 4)
|
|
):
|
|
disable_reason = "Only Deepseek V3/R1 on AMD-platform with capability >= gfx942(MI30x) can use shared experts fusion optimization under expert parallelism."
|
|
elif disable_reason is None and (
|
|
get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mori()
|
|
):
|
|
disable_reason = "Deepseek V3/R1 cannot use shared experts fusion optimization under deepep expert parallelism."
|
|
elif self.quant_config and self.quant_config.get_name() == "w4afp8":
|
|
disable_reason = "Deepseek V3/R1 W4AFP8 model uses different quant method for routed experts and shared experts."
|
|
|
|
if disable_reason is not None:
|
|
get_global_server_args().disable_shared_experts_fusion = True
|
|
self.num_fused_shared_experts = 0
|
|
log_info_on_rank0(
|
|
logger,
|
|
f"{disable_reason} Shared experts fusion optimization is disabled.",
|
|
)
|
|
return
|
|
|
|
self.num_fused_shared_experts = self.config.n_shared_experts
|
|
|
|
def get_input_embeddings(self) -> nn.Embedding:
|
|
return self.model.embed_tokens
|
|
|
|
@torch.no_grad()
|
|
def forward(
|
|
self,
|
|
input_ids: torch.Tensor,
|
|
positions: torch.Tensor,
|
|
forward_batch: ForwardBatch,
|
|
input_embeds: torch.Tensor = None,
|
|
pp_proxy_tensors: Optional[PPProxyTensors] = None,
|
|
) -> torch.Tensor:
|
|
if self.nsa_enable_prefill_cp:
|
|
if can_cp_split(len(input_ids), self.cp_size, self.use_nsa, forward_batch):
|
|
forward_batch.nsa_cp_metadata = prepare_input_dp_with_cp_dsa(
|
|
len(input_ids),
|
|
self.cp_rank,
|
|
self.cp_size,
|
|
forward_batch.seq_lens_cpu.tolist(),
|
|
forward_batch=forward_batch,
|
|
page_size=getattr(
|
|
getattr(forward_batch, "token_to_kv_pool", None),
|
|
"page_size",
|
|
None,
|
|
),
|
|
)
|
|
|
|
with get_attn_tp_context().maybe_input_scattered(forward_batch):
|
|
hidden_states = self.model(
|
|
input_ids, positions, forward_batch, input_embeds, pp_proxy_tensors
|
|
)
|
|
aux_hidden_states = None
|
|
if self.capture_aux_hidden_states:
|
|
hidden_states, aux_hidden_states = hidden_states
|
|
|
|
if self.pp_group.is_last_rank:
|
|
logits_output = self.logits_processor(
|
|
input_ids, hidden_states, self.lm_head, forward_batch, aux_hidden_states
|
|
)
|
|
logits_output.draft_hidden_states = getattr(
|
|
forward_batch, "draft_hidden_states", None
|
|
)
|
|
return logits_output
|
|
else:
|
|
return hidden_states
|
|
|
|
def _should_use_narrow_output_path(self, forward_batch: ForwardBatch) -> bool:
|
|
if not nsa_use_prefill_cp(forward_batch):
|
|
return False
|
|
if not self.pp_group.is_last_rank:
|
|
return False
|
|
if not forward_batch.forward_mode.is_extend():
|
|
return False
|
|
if forward_batch.return_logprob:
|
|
return False
|
|
if forward_batch.capture_hidden_mode != CaptureHiddenMode.NULL:
|
|
return False
|
|
return True
|
|
|
|
@property
|
|
def start_layer(self):
|
|
return self.model.start_layer
|
|
|
|
@property
|
|
def end_layer(self):
|
|
return self.model.end_layer
|
|
|
|
def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]], is_nextn=False):
|
|
self.do_load_weights(weights, is_nextn)
|
|
|
|
def get_embed_and_head(self):
|
|
return self.model.embed_tokens.weight, self.lm_head.weight
|
|
|
|
def set_embed_and_head(self, embed, head):
|
|
del self.model.embed_tokens.weight
|
|
del self.lm_head.weight
|
|
self.model.embed_tokens.weight = embed
|
|
self.lm_head.weight = head
|
|
torch.cuda.empty_cache()
|
|
torch.cuda.synchronize()
|
|
|
|
@classmethod
|
|
def get_model_config_for_expert_location(cls, config):
|
|
return ModelConfigForExpertLocation(
|
|
num_layers=config.num_hidden_layers,
|
|
num_logical_experts=config.n_routed_experts,
|
|
num_groups=config.n_group,
|
|
)
|
|
|
|
def set_eagle3_layers_to_capture(self, layer_ids: Optional[List[int]] = None):
|
|
if not self.pp_group.is_last_rank:
|
|
return
|
|
|
|
if layer_ids is None:
|
|
self.capture_aux_hidden_states = True
|
|
num_layers = self.config.num_hidden_layers
|
|
self.model.layers_to_capture = [2, num_layers // 2, num_layers - 3]
|
|
else:
|
|
self.capture_aux_hidden_states = True
|
|
# we plus 1 here because in sglang, for the ith layer, it takes the output
|
|
# of the (i-1)th layer as aux hidden state
|
|
self.model.layers_to_capture = [val + 1 for val in layer_ids]
|
|
|
|
|
|
class DeepseekV3ForCausalLM(DeepseekV2ForCausalLM):
|
|
pass
|
|
|
|
|
|
class DeepseekV32ForCausalLM(DeepseekV2ForCausalLM):
|
|
pass
|
|
|
|
|
|
EntryClass = [DeepseekV2ForCausalLM, DeepseekV3ForCausalLM, DeepseekV32ForCausalLM]
|