""" # Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved. # # 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. """ from __future__ import annotations import math import re from typing import Dict import paddle from paddle import nn from paddleformers.transformers import PretrainedModel from paddleformers.utils.log import logger from fastdeploy.config import FDConfig from fastdeploy.distributed.communication import tensor_model_parallel_all_reduce from fastdeploy.model_executor.forward_meta import ForwardMeta from fastdeploy.model_executor.graph_optimization.decorator import ( support_graph_optimization, ) from fastdeploy.model_executor.layers.activation import SiluAndMul from fastdeploy.model_executor.layers.attention.attention import Attention from fastdeploy.model_executor.layers.embeddings import VocabParallelEmbedding from fastdeploy.model_executor.layers.linear import ( ColumnParallelLinear, KVBatchLinear, MergedColumnParallelLinear, MergedReplicatedLinear, ReplicatedLinear, RowParallelLinear, ) from fastdeploy.model_executor.layers.lm_head import ParallelLMHead from fastdeploy.model_executor.layers.moe.moe import FusedMoE from fastdeploy.model_executor.layers.normalization import RMSNorm from fastdeploy.model_executor.layers.rotary_embedding import ( DeepseekScalingRotaryEmbedding, ) from fastdeploy.model_executor.models.model_base import ( ModelCategory, ModelForCasualLM, ModelRegistry, ) from fastdeploy.model_executor.ops.triton_ops.triton_utils import ( enable_compat_on_triton_kernel, ) from fastdeploy.platforms import current_platform if current_platform.is_cuda() or current_platform.is_maca(): from fastdeploy.model_executor.ops.gpu import ( get_position_ids_and_mask_encoder_batch, ) from fastdeploy.model_executor.layers.quantization.fp8_utils import ( per_token_group_quant_fp8, ) from fastdeploy.platforms import current_platform if current_platform.is_cuda(): from fastdeploy.model_executor.ops.gpu import ( cp_gather_indexer_k_quant_cache, indexer_k_quant_and_cache, radix_topk_ragged_transform, ) paddle.enable_compat(scope={"deep_gemm"}) class DeepSeekV3MLP(nn.Layer): """ DeepSeekV3MLP, for Dense FFN and Shared Experts Layer. """ def __init__( self, fd_config: FDConfig, intermediate_size: int, prefix: str = "", reduce_results: bool = True, ) -> None: super().__init__() self.up_gate_proj = MergedColumnParallelLinear( fd_config=fd_config, prefix=f"{prefix}.up_gate_proj", input_size=fd_config.model_config.hidden_size, output_size=intermediate_size * 2, with_bias=False, activation=fd_config.model_config.hidden_act, ) self.down_proj = RowParallelLinear( fd_config=fd_config, prefix=f"{prefix}.down_proj", input_size=intermediate_size, output_size=fd_config.model_config.hidden_size, with_bias=False, reduce_results=reduce_results, ) self.act_fn = SiluAndMul( fd_config=fd_config, bias=None, act_method=fd_config.model_config.hidden_act, ) def forward(self, x, forward_meta=None): """ """ gate_up_out = self.up_gate_proj(x) act_out = self.act_fn(gate_up_out) down_out = self.down_proj(act_out) return down_out class DeepSeekV3MoE(nn.Layer): """ DeepSeekV3MoE, for MoE Layer. """ def __init__(self, fd_config: FDConfig, layer_id: int, prefix: str) -> None: super().__init__() self.tp_size = fd_config.parallel_config.tensor_parallel_size self.ep_size = fd_config.parallel_config.expert_parallel_size self.attn_tp_size = fd_config.parallel_config.tensor_parallel_size if self.ep_size > 1: self.tp_size = 1 self.norm_topk_prob = fd_config.model_config.norm_topk_prob weight_key_map = { "gate_correction_bias_key": f"{prefix}.gate.e_score_correction_bias", "up_gate_proj_expert_weight_key": f"{prefix}.experts.{{}}.up_gate_proj.weight", "down_proj_expert_weight_key": f"{prefix}.experts.{{}}.down_proj.weight", } self.gate = ReplicatedLinear( fd_config=fd_config, prefix=f"{prefix}.gate", input_size=fd_config.model_config.hidden_size, output_size=fd_config.model_config.n_routed_experts, with_bias=False, skip_quant=True, weight_dtype="float32", ) if fd_config.model_config.topk_method == "noaux_tc": self.gate.e_score_correction_bias = self.create_parameter( shape=[1, fd_config.model_config.n_routed_experts], dtype="float32", default_initializer=paddle.nn.initializer.Constant(0), ) else: self.gate.e_score_correction_bias = None self.experts = FusedMoE( fd_config=fd_config, reduce_results=False, renormalize=self.norm_topk_prob, moe_intermediate_size=fd_config.model_config.moe_intermediate_size, num_experts=fd_config.model_config.n_routed_experts, top_k=fd_config.model_config.num_experts_per_tok, topk_method=fd_config.model_config.topk_method, topk_group=fd_config.model_config.topk_group, n_group=fd_config.model_config.n_group, routed_scaling_factor=fd_config.model_config.routed_scaling_factor, layer_idx=layer_id, gate_correction_bias=self.gate.e_score_correction_bias, weight_key_map=weight_key_map, ) self.num_shared_experts = fd_config.model_config.n_shared_experts shared_experts_intermediate_size = self.num_shared_experts * fd_config.model_config.moe_intermediate_size self.shared_experts = DeepSeekV3MLP( fd_config=fd_config, intermediate_size=shared_experts_intermediate_size, prefix=f"{prefix}.shared_experts", reduce_results=False, ) def forward(self, hidden_states: paddle.Tensor, forward_meta: ForwardMeta): """ """ shared_experts_out = self.shared_experts(hidden_states) if self.attn_tp_size > 1 and self.ep_size > 1: shared_experts_out = tensor_model_parallel_all_reduce(shared_experts_out) moe_out = self.experts(hidden_states, self.gate, forward_meta) moe_out = moe_out + shared_experts_out # We do to TP all reduce after the sum of experts. if self.tp_size > 1: moe_out = tensor_model_parallel_all_reduce(moe_out) return moe_out class DeepseekV3MLAAttention(nn.Layer): """ DeepseekV3MLAAttention """ def __init__(self, fd_config: FDConfig, layer_id: int, prefix: str = "") -> None: super().__init__() self.tp_size = fd_config.parallel_config.tensor_parallel_size self.hidden_size = fd_config.model_config.hidden_size self.num_attention_heads = fd_config.model_config.num_attention_heads self.num_attention_heads_tp = self.num_attention_heads // self.tp_size # MLA self.qk_nope_head_dim = fd_config.model_config.qk_nope_head_dim self.qk_rope_head_dim = fd_config.model_config.qk_rope_head_dim self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim self.v_head_dim = fd_config.model_config.v_head_dim self.q_lora_rank = fd_config.model_config.q_lora_rank self.kv_lora_rank = fd_config.model_config.kv_lora_rank self.attn_softmax_scale = self.qk_head_dim**-0.5 self.rope_theta = fd_config.model_config.rope_theta self.rms_norm_eps = fd_config.model_config.rms_norm_eps assert self.q_lora_rank is not None, "self.q_lora_rank is None, Please Check your config." # NOTE: (changwenbin) qkv_a_proj horizontal fusion self.qkv_a_proj_with_mqa = MergedReplicatedLinear( fd_config=fd_config, prefix=f"{prefix}.qkv_a_proj_with_mqa", input_size=self.hidden_size, output_sizes=[self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], with_bias=False, ) self.q_a_layernorm = RMSNorm( fd_config, hidden_size=self.q_lora_rank, eps=self.rms_norm_eps, prefix=f"{prefix}.q_a_layernorm", ) self.q_b_proj = ColumnParallelLinear( fd_config=fd_config, prefix=f"{prefix}.q_b_proj", input_size=self.q_lora_rank, output_size=self.num_attention_heads * self.qk_head_dim, with_bias=False, ) self.kv_a_layernorm = RMSNorm( fd_config, hidden_size=self.kv_lora_rank, eps=self.rms_norm_eps, prefix=f"{prefix}.kv_a_layernorm", ) self.kv_b_proj = ColumnParallelLinear( fd_config=fd_config, prefix=f"{prefix}.kv_b_proj", input_size=self.kv_lora_rank, output_size=self.num_attention_heads * (self.qk_nope_head_dim + self.v_head_dim), with_bias=False, ) self.o_proj = RowParallelLinear( fd_config, prefix=f"{prefix}.o_proj", input_size=self.num_attention_heads * self.v_head_dim, output_size=self.hidden_size, with_bias=False, layer_id=layer_id, ) self.kv_b_proj_bmm = KVBatchLinear( fd_config=fd_config, kv_b_proj=self.kv_b_proj, prefix=f"{prefix}.kv_b_proj", kv_lora_rank=self.kv_lora_rank, num_attention_heads=self.num_attention_heads, qk_nope_head_dim=self.qk_nope_head_dim, v_head_dim=self.v_head_dim, ) self.rope_scaling = getattr(fd_config.model_config, "rope_scaling", None) if self.rope_scaling: mscale_all_dim = self.rope_scaling.get("mscale_all_dim", False) scaling_factor = self.rope_scaling["factor"] mscale = self.yarn_get_mscale(scaling_factor, float(mscale_all_dim)) self.attn_softmax_scale = self.attn_softmax_scale * mscale * mscale rope_scaling_kwargs = { key: self.rope_scaling[key] for key in [ "beta_fast", "beta_slow", "mscale", "mscale_all_dim", ] if key in self.rope_scaling } self.rope_scaling_factor = self.rope_scaling["factor"] self.rope_scaling_original_max_position_embeddings = self.rope_scaling["original_max_position_embeddings"] self.rotary_emb = DeepseekScalingRotaryEmbedding( self.qk_rope_head_dim, max_position_embeddings=self.rope_scaling_original_max_position_embeddings, base=self.rope_theta, scaling_factor=self.rope_scaling_factor, **rope_scaling_kwargs, ) else: # Default rope without scaling max_position_embeddings = getattr(fd_config.model_config, "max_position_embeddings", 8192) self.rotary_emb = DeepseekScalingRotaryEmbedding( self.qk_rope_head_dim, max_position_embeddings=max_position_embeddings, base=self.rope_theta, scaling_factor=1.0, ) self.mla_attn = Attention( fd_config=fd_config, layer_id=layer_id, prefix=prefix, use_neox_rotary_style=False, ) self.prefix = prefix @staticmethod def yarn_get_mscale(scale=1, mscale=1): """ """ if scale <= 1: return 1.0 return 0.1 * mscale * math.log(scale) + 1.0 def forward( self, forward_meta: ForwardMeta, hidden_states: paddle.Tensor, position_ids: paddle.Tensor, mask_encoder_batch: paddle.Tensor, ): """ """ fmha_out = None # NOTE: (changwenbin) qkv_a_proj horizontal fusion qkv_a_out = self.qkv_a_proj_with_mqa(hidden_states) query, compressed_kv, key_pe = qkv_a_out.split( [self.q_lora_rank, self.kv_lora_rank, self.qk_rope_head_dim], axis=-1 ) query = self.q_a_layernorm(query)[0] query = self.q_b_proj(query) query.reshape_([-1, self.num_attention_heads_tp, self.qk_head_dim]) query_nope, query_pe = query.split([self.qk_nope_head_dim, self.qk_rope_head_dim], axis=-1) key_pe.reshape_([-1, 1, self.qk_rope_head_dim]) query_pe, key_pe = self.rotary_emb(position_ids, query_pe, key_pe) compressed_kv = self.kv_a_layernorm(compressed_kv)[0] need_do_prefill = forward_meta.max_len_tensor_cpu[1] > 0 need_do_decode = forward_meta.max_len_tensor_cpu[2] > 0 if need_do_prefill: # max_enc_len_this_time key_value = self.kv_b_proj(compressed_kv) key_value.reshape_( [ -1, self.num_attention_heads_tp, self.qk_nope_head_dim + self.v_head_dim, ] ) key_nope, value = key_value.split([self.qk_nope_head_dim, self.v_head_dim], axis=-1) query[..., self.qk_nope_head_dim :] = query_pe key = paddle.empty_like(query) key[..., : self.qk_nope_head_dim] = key_nope key[..., self.qk_nope_head_dim :] = key_pe value = paddle.nn.functional.pad(value, [0, self.qk_head_dim - self.v_head_dim], value=0) fmha_out_prefill = self.mla_attn( q=query, k=key, v=value, qkv=None, compressed_kv=compressed_kv, k_pe=key_pe, forward_meta=forward_meta, ) fmha_out_prefill.reshape_([-1, self.num_attention_heads_tp, self.qk_head_dim]) fmha_out_prefill = fmha_out_prefill[:, :, : self.v_head_dim] fmha_out_prefill.reshape_([-1, self.num_attention_heads_tp * self.v_head_dim]) fmha_out_prefill = fmha_out_prefill * mask_encoder_batch.cast(fmha_out_prefill.dtype) fmha_out = fmha_out_prefill if need_do_decode: # max_dec_len_this_time q_nope_out = self.kv_b_proj_bmm(query_nope.transpose([1, 0, 2]), proj_type="k").transpose([1, 0, 2]) q_input = paddle.concat([q_nope_out, query_pe], axis=-1) q_input.reshape_( [ -1, self.num_attention_heads_tp * (self.kv_lora_rank + self.qk_rope_head_dim), ] ) fmha_out_decode = self.mla_attn( q=q_input, k=None, v=None, qkv=None, compressed_kv=compressed_kv, k_pe=key_pe, forward_meta=forward_meta, ) fmha_out_decode = fmha_out_decode.reshape_([-1, self.num_attention_heads_tp, self.kv_lora_rank]).transpose( [1, 0, 2] ) fmha_out_decode = ( self.kv_b_proj_bmm(fmha_out_decode, proj_type="v") .transpose([1, 0, 2]) .reshape_([-1, self.num_attention_heads_tp * self.v_head_dim]) ) if need_do_prefill: fmha_out += fmha_out_decode else: fmha_out = fmha_out_decode output = self.o_proj(fmha_out) return output def compute_slot_mapping( block_tables: paddle.Tensor, # [num_reqs, max_blocks_per_req] positions: paddle.Tensor, # [num_tokens] 每个token的位置 batch_id_per_token: paddle.Tensor, # [num_tokens] 每个token属于哪个请求 block_size: int, ) -> paddle.Tensor: """ 计算 slot_mapping 公式: slot = block_id * block_size + offset_in_block """ # 1. 计算每个 token 对应的 block 索引 block_idx = positions // block_size # [num_tokens] # 2. 从 block_tables 中查表获取 block_id # block_tables[batch_id_per_token, block_idx] block_ids = block_tables[batch_id_per_token, block_idx] # [num_tokens] # 3. 计算在 block 内的偏移 block_offset = positions % block_size # [num_tokens] # 4. 计算 slot_mapping slot_mapping = block_ids * block_size + block_offset return slot_mapping.cast(paddle.int64) import triton import triton.language as tl @enable_compat_on_triton_kernel @triton.jit() def extract_kernel( q, weight, cu_seqlens_q, seq_lens_encoder, seq_lens_decoder, output, out_weight, cache_seqlens, HIDDEN_DIM: tl.constexpr, WEIGHT_DIM: tl.constexpr, BLOCK_SIZE: tl.constexpr, ): batch_id = tl.program_id(axis=0) cache_kv_len = tl.load(seq_lens_decoder + batch_id) # 这个batch不是decoder，所以不需要动弹 if cache_kv_len <= 0: return cu_len_this_batch = tl.load(cu_seqlens_q + batch_id) read_offsets = tl.arange(0, BLOCK_SIZE) read_weight_offsets = tl.arange(0, WEIGHT_DIM) q += cu_len_this_batch * HIDDEN_DIM weight += cu_len_this_batch * WEIGHT_DIM row_data = tl.load(q + read_offsets, mask=read_offsets < HIDDEN_DIM) weight_row_data = tl.load(weight + read_weight_offsets, mask=read_weight_offsets < WEIGHT_DIM) output += batch_id * HIDDEN_DIM out_weight += batch_id * WEIGHT_DIM tl.store(output + read_offsets, row_data, mask=read_offsets < HIDDEN_DIM) tl.store(out_weight + read_weight_offsets, weight_row_data, mask=read_weight_offsets < WEIGHT_DIM) tl.store(cache_seqlens + batch_id, cache_kv_len + 1) def extract_decoder_token_from_q( q: paddle.Tensor, weight: paddle.Tensor, cu_seqlens_q: paddle.Tensor, seq_lens_encoder: paddle.Tensor, seq_lens_decoder: paddle.Tensor, ): assert len(q.shape) == 2 assert len(weight.shape) == 2 assert len(cu_seqlens_q.shape) == 1 assert len(seq_lens_encoder.shape) == 1 assert len(seq_lens_decoder.shape) == 1 max_bsz = seq_lens_decoder.shape[0] hidden_dim = q.shape[-1] weight_dim = weight.shape[-1] # if q.shape[0] <= max_bsz: # max_bsz = q.shape[0] out = paddle.zeros([max_bsz, hidden_dim], dtype=q.dtype) out_weight = paddle.zeros([max_bsz, weight_dim], dtype=weight.dtype) cache_seqlens = paddle.zeros_like(seq_lens_decoder) BLOCK_SIZE = triton.next_power_of_2(hidden_dim) grid = (max_bsz,) extract_kernel[grid]( q, weight, cu_seqlens_q, seq_lens_encoder, seq_lens_decoder, out, out_weight, cache_seqlens, hidden_dim, weight_dim, BLOCK_SIZE, ) return out, out_weight, cache_seqlens class Indexer(nn.Layer): def __init__( self, fd_config: FDConfig, layer_id: int, prefix: str = "", ): super().__init__() self.layer_id = layer_id self.max_model_len = fd_config.model_config.max_model_len self.index_head_dim = fd_config.model_config.index_head_dim self.index_n_heads = fd_config.model_config.index_n_heads self.index_topk = fd_config.model_config.index_topk self.rope_dim = fd_config.model_config.qk_rope_head_dim # 64 self.q_lora_rank = fd_config.model_config.q_lora_rank # 1536 self.hidden_size = fd_config.model_config.hidden_size self.wq_b = ReplicatedLinear( fd_config=fd_config, prefix=f"{prefix}.wq_b", input_size=self.q_lora_rank, output_size=self.index_head_dim * self.index_n_heads, with_bias=False, ) self.wk = ReplicatedLinear( fd_config=fd_config, prefix=f"{prefix}.wk", input_size=self.hidden_size, output_size=self.index_head_dim, with_bias=False, ) self.k_norm = RMSNorm(fd_config, self.index_head_dim, eps=1e-6, prefix=f"{prefix}.k_norm") # self.k_norm = LayerNorm(self.head_dim, eps=1e-6) self.weights_proj = ReplicatedLinear( fd_config=fd_config, prefix=f"{prefix}.weights_proj", input_size=self.hidden_size, output_size=self.index_n_heads, with_bias=False, ) self.softmax_scale = self.index_head_dim**-0.5 self.scale_fmt = "ue8m0" self.quant_block_size = 128 # TODO: get from config self.offsets = paddle.zeros([self.max_model_len], dtype="int32") self.lengths = paddle.zeros([self.max_model_len], dtype="int32") # self.buffer = paddle.zeros([2048 * 2048], dtype=paddle.uint8) def forward( self, forward_meta: ForwardMeta, hidden_states: paddle.Tensor, qr: paddle.Tensor, positions, rotary_emb ) -> paddle.Tensor: self.indexer_cache = forward_meta.caches[2 * self.layer_id + 1] q = self.wq_b(qr) q = q.reshape([-1, self.index_n_heads, self.index_head_dim]) q_pe, q_nope = paddle.split(q, [self.rope_dim, self.index_head_dim - self.rope_dim], axis=-1) k = self.wk(hidden_states) k, _ = self.k_norm(k) k_pe, k_nope = paddle.split(k, [self.rope_dim, self.index_head_dim - self.rope_dim], axis=-1) q_pe, k_pe = rotary_emb(positions, q_pe, k_pe.unsqueeze(1)) q_pe = q_pe.reshape(-1, self.index_n_heads, self.rope_dim) k_pe = k_pe.reshape(-1, 1, self.rope_dim) # [num_tokens, n_head, rope_dim]. q = paddle.concat([q_pe, q_nope], axis=-1).reshape([-1, self.index_head_dim]) # `k_pe` is [num_tokens, 1, rope_dim] (MQA). k = paddle.concat([k_pe.squeeze(-2), k_nope], axis=-1) # indexer q_quant q_fp8, q_scale = per_token_group_quant_fp8( q, self.quant_block_size, column_major_scales=False, use_ue8m0=self.scale_fmt is not None, ) q_fp8 = q_fp8.reshape([-1, self.index_n_heads, self.index_head_dim]) q_scale = q_scale.reshape([-1, self.index_n_heads, 1]) weights = self.weights_proj(hidden_states) weights = weights.unsqueeze(-1) * q_scale * self.softmax_scale * self.index_n_heads**-0.5 weights = weights.squeeze(-1) slot_mapping = compute_slot_mapping( forward_meta.block_tables, forward_meta.position_ids, forward_meta.batch_id_per_token, 64, ) indexer_top_k = paddle.full([q_fp8.shape[0], self.index_topk], -1, dtype="int32") # indexer write_cache indexer_k_quant_and_cache(k, self.indexer_cache, slot_mapping, self.quant_block_size, self.scale_fmt) import deep_gemm if forward_meta.max_len_tensor_cpu[1]: # indexer_prefill read_cache k_fp8_cache = paddle.zeros_like(k, dtype=paddle.uint8) k_scale_cache = paddle.zeros([k.shape[0], 4], dtype=paddle.float32) cp_gather_indexer_k_quant_cache( self.indexer_cache, k_fp8_cache, k_scale_cache, forward_meta.block_tables, forward_meta.cu_seqlens_k ) k_scale_cache = k_scale_cache.flatten()[: k.shape[0]] k_cache = k_fp8_cache.view(paddle.float8_e4m3fn), k_scale_cache # TODO(changwenbin): Constructed using maskoffset # ks,ke = forward_meta.attn_mask_offsets[::2].contiguous(),forward_meta.attn_mask_offsets[1::2].contiguous() num_tokens = q_fp8.shape[0] ks = paddle.zeros(num_tokens, dtype=paddle.int32) ks_topk = paddle.zeros(num_tokens, dtype=paddle.int32) ke = paddle.zeros(num_tokens, dtype=paddle.int32) bsz = forward_meta.seq_lens_this_time.shape[0] for i in range(bsz): if forward_meta.seq_lens_encoder[i] > 0: token_start_k = forward_meta.cu_seqlens_k[i] token_end_k = forward_meta.cu_seqlens_k[i + 1] ks[token_start_k:token_end_k] = forward_meta.cu_seqlens_k[i] ke[token_start_k:token_end_k] = paddle.arange(token_start_k, token_end_k, dtype=paddle.int32) + 1 max_seqlen_k = (ke - ks).max().item() logits = deep_gemm.fp8_mqa_logits( q_fp8, k_cache, weights, ks, ke, max_seqlen_k=max_seqlen_k, clean_logits=False ) # To save GPU global memory usage assert logits.size() == (num_tokens, max_seqlen_k) tmp = paddle.full((num_tokens, num_tokens), float("-inf")) for i in range(num_tokens): tmp[i, ks[i] : ke[i]] = logits[i, : ke[i] - ks[i]] logits = tmp radix_topk_ragged_transform( logits.contiguous(), indexer_top_k, ks_topk, # self.offsets, ke - ks + 1, # mask.contiguous(),#self.lengths, None, # forward_meta.seq_lens_decoder, None, # forward_meta.batch_id_per_token, None, None, # self.buffer 0, self.index_topk, 1, ) if forward_meta.max_len_tensor_cpu[2]: decoder_q, decoder_weight, cache_seqlens = extract_decoder_token_from_q( q_fp8.reshape(-1, self.index_n_heads * self.index_head_dim), weights, forward_meta.cu_seqlens_q, forward_meta.seq_lens_encoder, forward_meta.seq_lens_decoder, ) schedule_metadata = deep_gemm.get_paged_mqa_logits_metadata(cache_seqlens, 64, deep_gemm.get_num_sms()) logits = deep_gemm.fp8_paged_mqa_logits( decoder_q.reshape(-1, 1, self.index_n_heads, self.index_head_dim), self.indexer_cache.unsqueeze(2), decoder_weight, cache_seqlens, forward_meta.block_tables, schedule_metadata, self.max_model_len, clean_logits=True, ) radix_topk_ragged_transform( logits.contiguous(), indexer_top_k, self.offsets, # unused self.lengths, # unused cache_seqlens, forward_meta.batch_id_per_token, forward_meta.block_tables, None, # self.buffer forward_meta.block_tables.shape[1], self.index_topk, 1, # q_head ) return indexer_top_k class DeepseekV32DSAAttention(nn.Layer): """ DeepseekV32DSAAttention """ def __init__(self, fd_config: FDConfig, layer_id: int, prefix: str = "") -> None: super().__init__() self.tp_size = fd_config.parallel_config.tensor_parallel_size self.hidden_size = fd_config.model_config.hidden_size self.num_attention_heads = fd_config.model_config.num_attention_heads self.num_attention_heads_tp = self.num_attention_heads // self.tp_size # MLA self.qk_nope_head_dim = fd_config.model_config.qk_nope_head_dim self.qk_rope_head_dim = fd_config.model_config.qk_rope_head_dim self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim self.v_head_dim = fd_config.model_config.v_head_dim self.q_lora_rank = fd_config.model_config.q_lora_rank self.kv_lora_rank = fd_config.model_config.kv_lora_rank # Indexer self.index_head_dim = fd_config.model_config.index_head_dim self.index_n_heads = fd_config.model_config.index_n_heads self.index_topk = fd_config.model_config.index_topk self.attn_softmax_scale = self.qk_head_dim**-0.5 self.rope_theta = fd_config.model_config.rope_theta if fd_config.model_config.model_type == "glm_moe_dsa": self.rope_theta = fd_config.model_config.rope_parameters["rope_theta"] self.rms_norm_eps = fd_config.model_config.rms_norm_eps assert self.q_lora_rank is not None, "self.q_lora_rank is None, Please Check your config." self.qkv_a_proj_with_mqa = MergedReplicatedLinear( fd_config=fd_config, prefix=f"{prefix}.qkv_a_proj_with_mqa", input_size=self.hidden_size, output_sizes=[self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], with_bias=False, ) self.q_a_layernorm = RMSNorm( fd_config, hidden_size=self.q_lora_rank, eps=self.rms_norm_eps, prefix=f"{prefix}.q_a_layernorm", ) self.q_b_proj = ColumnParallelLinear( fd_config=fd_config, prefix=f"{prefix}.q_b_proj", input_size=self.q_lora_rank, output_size=self.num_attention_heads * self.qk_head_dim, with_bias=False, ) self.kv_a_layernorm = RMSNorm( fd_config, hidden_size=self.kv_lora_rank, eps=self.rms_norm_eps, prefix=f"{prefix}.kv_a_layernorm", ) self.kv_b_proj = ColumnParallelLinear( fd_config=fd_config, prefix=f"{prefix}.kv_b_proj", input_size=self.kv_lora_rank, output_size=self.num_attention_heads * (self.qk_nope_head_dim + self.v_head_dim), with_bias=False, ) self.o_proj = RowParallelLinear( fd_config, prefix=f"{prefix}.o_proj", input_size=self.num_attention_heads * self.v_head_dim, output_size=self.hidden_size, with_bias=False, layer_id=layer_id, ) self.kv_b_proj_bmm = KVBatchLinear( fd_config=fd_config, kv_b_proj=self.kv_b_proj, prefix=f"{prefix}.kv_b_proj", kv_lora_rank=self.kv_lora_rank, num_attention_heads=self.num_attention_heads, qk_nope_head_dim=self.qk_nope_head_dim, v_head_dim=self.v_head_dim, ) self.rope_scaling = getattr(fd_config.model_config, "rope_scaling", None) if self.rope_scaling: mscale_all_dim = self.rope_scaling.get("mscale_all_dim", False) scaling_factor = self.rope_scaling["factor"] mscale = self.yarn_get_mscale(scaling_factor, float(mscale_all_dim)) self.attn_softmax_scale = self.attn_softmax_scale * mscale * mscale rope_scaling_kwargs = { key: self.rope_scaling[key] for key in [ "beta_fast", "beta_slow", "mscale", "mscale_all_dim", ] if key in self.rope_scaling } self.rope_scaling_factor = self.rope_scaling["factor"] self.rope_scaling_original_max_position_embeddings = self.rope_scaling["original_max_position_embeddings"] self.rotary_emb = DeepseekScalingRotaryEmbedding( self.qk_rope_head_dim, max_position_embeddings=self.rope_scaling_original_max_position_embeddings, base=self.rope_theta, scaling_factor=self.rope_scaling_factor, **rope_scaling_kwargs, ) self.indexer_rotary_emb = DeepseekScalingRotaryEmbedding( self.qk_rope_head_dim, max_position_embeddings=self.rope_scaling_original_max_position_embeddings, base=self.rope_theta, scaling_factor=self.rope_scaling_factor, **rope_scaling_kwargs, ) else: # Default rope without scaling max_position_embeddings = getattr(fd_config.model_config, "max_position_embeddings", 8192) self.rotary_emb = DeepseekScalingRotaryEmbedding( self.qk_rope_head_dim, max_position_embeddings=max_position_embeddings, base=self.rope_theta, scaling_factor=1.0, ) self.indexer_rotary_emb = DeepseekScalingRotaryEmbedding( self.qk_rope_head_dim, max_position_embeddings=max_position_embeddings, base=self.rope_theta, scaling_factor=1.0, ) self.indexer = Indexer( fd_config=fd_config, layer_id=layer_id, prefix=prefix, ) self.dsa_attn = Attention( fd_config=fd_config, layer_id=layer_id, prefix=prefix, use_neox_rotary_style=False, ) self.prefix = prefix @staticmethod def yarn_get_mscale(scale=1, mscale=1): """ """ if scale <= 1: return 1.0 return 0.1 * mscale * math.log(scale) + 1.0 def forward( self, forward_meta: ForwardMeta, hidden_states: paddle.Tensor, position_ids: paddle.Tensor, mask_encoder_batch: paddle.Tensor, ): """ """ qkv_a_out = self.qkv_a_proj_with_mqa(hidden_states) query, compressed_kv, key_pe = qkv_a_out.split( [self.q_lora_rank, self.kv_lora_rank, self.qk_rope_head_dim], axis=-1 ) key_pe.reshape_([-1, 1, self.qk_rope_head_dim]) query = self.q_a_layernorm(query)[0] # DSA indexer indexer_top_k = self.indexer( forward_meta, hidden_states, query, position_ids, rotary_emb=self.indexer_rotary_emb ) query = self.q_b_proj(query) query.reshape_([-1, self.num_attention_heads_tp, self.qk_head_dim]) query_nope, query_pe = query.split([self.qk_nope_head_dim, self.qk_rope_head_dim], axis=-1) query_pe, key_pe = self.rotary_emb(position_ids, query_pe, key_pe) q_nope_out = self.kv_b_proj_bmm(query_nope.transpose([1, 0, 2]).contiguous(), proj_type="k") q_input = paddle.concat([q_nope_out.transpose([1, 0, 2]).contiguous(), query_pe], axis=-1) compressed_kv = self.kv_a_layernorm(compressed_kv)[0] kv = paddle.concat([compressed_kv, key_pe.squeeze(1)], axis=-1) # dsa attention fmha_out = self.dsa_attn( q=q_input.contiguous(), k=kv.unsqueeze(1).contiguous(), v=indexer_top_k.unsqueeze(1).contiguous(), qkv=None, compressed_kv=compressed_kv, k_pe=key_pe, forward_meta=forward_meta, ) fmha_out = fmha_out.reshape_([-1, self.num_attention_heads_tp, self.kv_lora_rank]).transpose([1, 0, 2]) fmha_out = ( self.kv_b_proj_bmm( fmha_out, proj_type="v", ) .transpose([1, 0, 2]) .reshape_([-1, self.num_attention_heads_tp * self.v_head_dim]) ) output = self.o_proj(fmha_out) return output class DeepSeekV3DecoderLayer(nn.Layer): """ DeepSeekV3DecoderLayer """ def __init__( self, fd_config: FDConfig, prefix: str = "", ) -> None: super().__init__() layer_id = int(prefix.split(sep=".")[-1]) if fd_config.model_config.model_type in ["deepseek_v32", "glm_moe_dsa"]: self.self_attn = DeepseekV32DSAAttention( fd_config=fd_config, layer_id=layer_id, prefix=f"{prefix}.self_attn", ) else: self.self_attn = DeepseekV3MLAAttention( fd_config=fd_config, layer_id=layer_id, prefix=f"{prefix}.self_attn", ) if ( fd_config.model_config.n_routed_experts is not None and layer_id >= fd_config.model_config.first_k_dense_replace ): self.mlp = DeepSeekV3MoE( fd_config=fd_config, layer_id=layer_id, prefix=f"{prefix}.mlp", ) else: self.mlp = DeepSeekV3MLP( fd_config=fd_config, intermediate_size=fd_config.model_config.intermediate_size, prefix=f"{prefix}.mlp", ) self.input_layernorm = RMSNorm( fd_config, hidden_size=fd_config.model_config.hidden_size, eps=fd_config.model_config.rms_norm_eps, prefix=f"{prefix}.input_layernorm", layer_id=layer_id, ) self.post_attention_layernorm = RMSNorm( fd_config, hidden_size=fd_config.model_config.hidden_size, eps=fd_config.model_config.rms_norm_eps, prefix=f"{prefix}.post_attention_layernorm", layer_id=layer_id, ) def forward( self, forward_meta: ForwardMeta, hidden_states: paddle.Tensor, residual: paddle.Tensor, position_ids: paddle.Tensor, mask_encoder_batch: paddle.Tensor, ): """ """ if hidden_states.shape[0] > 0: hidden_states, residual = self.input_layernorm( hidden_states, residual_input=residual, forward_meta=forward_meta ) hidden_states = self.self_attn(forward_meta, hidden_states, position_ids, mask_encoder_batch) hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) else: residual = hidden_states hidden_states = self.mlp(hidden_states, forward_meta) return hidden_states, residual @support_graph_optimization class DeepSeekV3Model(nn.Layer): """ DeepSeekV3Model """ def __init__( self, fd_config: FDConfig = None, ): """ Initializer for the DeepSeekV3Model class. """ super().__init__() self.num_layers = fd_config.model_config.num_hidden_layers fd_config.model_config.pretrained_config.prefix_name = "deepseek_v3" self.embed_tokens = VocabParallelEmbedding( fd_config, num_embeddings=fd_config.model_config.vocab_size, embedding_dim=fd_config.model_config.hidden_size, params_dtype=paddle.get_default_dtype(), prefix="deepseek_v3.embed_tokens", ) self.layers = nn.LayerList( [ DeepSeekV3DecoderLayer( fd_config, prefix=f"{fd_config.model_config.pretrained_config.prefix_name}.layers.{i}", ) for i in range(self.num_layers) ] ) self.norm = RMSNorm( fd_config, hidden_size=fd_config.model_config.hidden_size, eps=fd_config.model_config.rms_norm_eps, prefix="deepseek_v3.norm", ) def forward( self, ids_remove_padding: paddle.Tensor, forward_meta: ForwardMeta, position_ids: paddle.Tensor, mask_encoder_batch: paddle.Tensor, ): """ """ hidden_states = self.embed_tokens(ids_remove_padding=ids_remove_padding, forward_meta=forward_meta) residual = None for i in range(self.num_layers): hidden_states, residual = self.layers[i]( forward_meta, hidden_states, residual, position_ids, mask_encoder_batch, ) out = self.norm(hidden_states, residual, forward_meta=forward_meta)[0] if self.norm.is_last_norm and self.norm.fd_config.parallel_config.use_sequence_parallel_moe: out = self.norm.allgather(out, forward_meta.ids_remove_padding.shape[0]) return out @ModelRegistry.register_model_class( architecture="DeepseekV3ForCausalLM", module_name="deepseek_v3", category=ModelCategory.TEXT_GENERATION, primary_use=ModelCategory.TEXT_GENERATION, ) class DeepseekV3ForCausalLM(ModelForCasualLM): """ DeepseekV3ForCausalLM """ def __init__(self, fd_config: FDConfig): """ Args: fd_config (FDConfig): Configurations for the LLM model. """ super().__init__(fd_config) self.model = DeepSeekV3Model(fd_config) self.ori_vocab_size = fd_config.model_config.ori_vocab_size self.lm_head = ParallelLMHead( fd_config, embedding_dim=fd_config.model_config.hidden_size, num_embeddings=fd_config.model_config.vocab_size, prefix="lm_head", ) self.position_ids_buffer = paddle.empty( [fd_config.scheduler_config.max_num_batched_tokens], dtype=paddle.int32 ) self.mask_encoder_batch_buffer = paddle.empty( [fd_config.scheduler_config.max_num_batched_tokens, 1], dtype=paddle.int32 ) @classmethod def name(cls): """ """ return "DeepseekV3ForCausalLM" @paddle.no_grad() def set_state_dict(self, state_dict): """ Load model parameters from a given state dictionary. """ pass @paddle.no_grad() def load_weights(self, weights_iterator) -> None: """ Load model parameters from a given weights_iterator object. Args: weights_iterator (Iterator): An iterator yielding (name, weight) pairs. """ from fastdeploy.model_executor.utils import ( default_weight_loader, process_weights_after_loading, ) stacked_params_mapping = [ # (param_name, shard_name, shard_id) ("up_gate_proj", "gate_proj", "gate"), ("up_gate_proj", "up_proj", "up"), ("embed_tokens.embeddings", "embed_tokens", None), ("lm_head.linear", "lm_head", None), ("experts.gate_correction_bias", "gate.e_score_correction_bias", None), ("qkv_a_proj_with_mqa", "q_a_proj", "q_a"), ("qkv_a_proj_with_mqa", "kv_a_proj_with_mqa", "kv_a"), ] # (param_name, weight_name, expert_id, shard_id) expert_params_mapping = FusedMoE.make_expert_params_mapping( num_experts=self.fd_config.model_config.n_routed_experts, ckpt_gate_proj_name="gate_proj", ckpt_down_proj_name="down_proj", ckpt_up_proj_name="up_proj", param_gate_up_proj_name="experts.up_gate_proj_", param_down_proj_name="experts.down_proj_", ) params_dict = dict(self.named_parameters()) process_weights_after_loading_fn = process_weights_after_loading(dict(self.named_sublayers()), self.fd_config) for loaded_weight_name, loaded_weight in weights_iterator: logger.debug(f"Loading weight: {loaded_weight_name}") loaded_weight_name = loaded_weight_name.replace("deepseek_v3", "model") for param_name, weight_name, shard_id in stacked_params_mapping: if weight_name not in loaded_weight_name: continue if "mlp.experts." in loaded_weight_name: continue model_param_name = loaded_weight_name.replace(weight_name, param_name) if model_param_name not in params_dict: continue param = params_dict[model_param_name] weight_loader = getattr(param, "weight_loader", default_weight_loader(self.fd_config)) weight_loader(param, loaded_weight, shard_id) break else: for mapping in expert_params_mapping: param_name, weight_name, expert_id, shard_id = mapping if weight_name not in loaded_weight_name: continue model_param_name = loaded_weight_name.replace(weight_name, param_name) if model_param_name not in params_dict: continue param = params_dict[model_param_name] weight_loader = param.weight_loader weight_loader(param, loaded_weight, shard_id=shard_id, expert_id=expert_id) break else: model_param_name = loaded_weight_name if model_param_name not in params_dict: continue param = params_dict[model_param_name] weight_loader = getattr(param, "weight_loader", default_weight_loader(self.fd_config)) weight_loader(param, loaded_weight) model_sublayer_name = re.sub(r"\.(up_gate_proj_weight|down_proj_weight|weight)$", "", model_param_name) if "kv_b_proj" in model_sublayer_name: kv_model_sublayer_name = model_sublayer_name.replace("kv_b_proj", "kv_b_proj_bmm") process_weights_after_loading_fn(kv_model_sublayer_name) process_weights_after_loading_fn(model_sublayer_name, param) def compute_logits(self, hidden_states: paddle.Tensor, forward_meta: ForwardMeta = None): """ """ logits = self.lm_head(hidden_states) logits = logits.astype(paddle.float32) logits[:, self.ori_vocab_size :] = -float("inf") return logits def pre_process(self, forward_meta): """ """ seq_lens_encoder = forward_meta.seq_lens_encoder seq_lens_decoder = forward_meta.seq_lens_decoder seq_lens_this_time = forward_meta.seq_lens_this_time current_total_tokens = forward_meta.ids_remove_padding.shape[0] position_ids = self.position_ids_buffer[:current_total_tokens] mask_encoder_batch = self.mask_encoder_batch_buffer[:current_total_tokens] get_position_ids_and_mask_encoder_batch( seq_lens_encoder, seq_lens_decoder, seq_lens_this_time, position_ids, mask_encoder_batch, ) return position_ids, mask_encoder_batch def empty_input_forward(self, forward_meta): """ empty_input_forward """ fake_hidden_states = paddle.empty( shape=[1, self.fd_config.model_config.hidden_size], dtype=paddle.get_default_dtype(), ) for i in range( self.fd_config.model_config.first_k_dense_replace, self.fd_config.model_config.num_hidden_layers, ): self.model.layers[i].mlp.experts(fake_hidden_states, self.model.layers[i].mlp.gate, forward_meta) def forward( self, inputs: Dict, forward_meta: ForwardMeta, ): ids_remove_padding = inputs["ids_remove_padding"] forward_meta.position_ids, mask_encoder_batch = self.pre_process(forward_meta) hidden_states = self.model( ids_remove_padding=ids_remove_padding, forward_meta=forward_meta, position_ids=forward_meta.position_ids, mask_encoder_batch=mask_encoder_batch, ) return hidden_states def clear_grpah_opt_backend(self): """Clear graph optimization backend, the captured cuda graph will be cleaned""" self.model.clear_grpah_opt_backend(fd_config=self.fd_config) class DeepSeekV3PretrainedModel(PretrainedModel): """ DeepSeekV3PretrainedModel """ config_class = FDConfig def _init_weight(self, layer): """ _init_weight """ return None @classmethod def arch_name(self): return "DeepseekV3ForCausalLM" @ModelRegistry.register_model_class( architecture="DeepseekV32ForCausalLM", module_name="deepseek_v3", # TODO(changwenbin): trick using the current dsk-v3 model category=ModelCategory.TEXT_GENERATION, primary_use=ModelCategory.TEXT_GENERATION, ) class DeepseekV32ForCausalLM(DeepseekV3ForCausalLM): """ DeepseekV32ForCausalLM """ @classmethod def name(cls): return "DeepseekV32ForCausalLM" class DeepSeekV32PretrainedModel(DeepSeekV3PretrainedModel): """ DeepSeekV32PretrainedModel """ @classmethod def arch_name(self): return "DeepseekV32ForCausalLM"