diff --git a/lightllm/common/basemodel/attention/create_utils.py b/lightllm/common/basemodel/attention/create_utils.py
index 594e81a9b4..95493d1de2 100644
--- a/lightllm/common/basemodel/attention/create_utils.py
+++ b/lightllm/common/basemodel/attention/create_utils.py
@@ -130,21 +130,25 @@ def get_mla_decode_att_backend_class(index=0, priority_list: list = ["flashinfer
return _auto_select_backend(llm_dtype, kv_type_to_backend=mla_data_type_to_backend, priority_list=priority_list)
-def get_nsa_prefill_att_backend_class(index=0, priority_list: list = ["flashmla_sparse"]) -> BaseAttBackend:
+def get_nsa_prefill_att_backend_class(
+ index=0, priority_list: list = ["flashmla_sparse"], backend_map=nsa_data_type_to_backend
+) -> BaseAttBackend:
args = get_env_start_args()
llm_dtype = args.llm_kv_type
backend_str = args.llm_prefill_att_backend[index]
if backend_str != "auto":
- return nsa_data_type_to_backend[llm_dtype][backend_str]
+ return backend_map[llm_dtype][backend_str]
else:
- return _auto_select_backend(llm_dtype, kv_type_to_backend=nsa_data_type_to_backend, priority_list=priority_list)
+ return _auto_select_backend(llm_dtype, kv_type_to_backend=backend_map, priority_list=priority_list)
-def get_nsa_decode_att_backend_class(index=0, priority_list: list = ["flashmla_sparse"]) -> BaseAttBackend:
+def get_nsa_decode_att_backend_class(
+ index=0, priority_list: list = ["flashmla_sparse"], backend_map=nsa_data_type_to_backend
+) -> BaseAttBackend:
args = get_env_start_args()
llm_dtype = args.llm_kv_type
backend_str = args.llm_decode_att_backend[index]
if backend_str != "auto":
- return nsa_data_type_to_backend[llm_dtype][backend_str]
+ return backend_map[llm_dtype][backend_str]
else:
- return _auto_select_backend(llm_dtype, kv_type_to_backend=nsa_data_type_to_backend, priority_list=priority_list)
+ return _auto_select_backend(llm_dtype, kv_type_to_backend=backend_map, priority_list=priority_list)
diff --git a/lightllm/common/basemodel/attention/nsa/dsv4_fp8_flashmla_sparse.py b/lightllm/common/basemodel/attention/nsa/dsv4_fp8_flashmla_sparse.py
new file mode 100644
index 0000000000..eece2648fb
--- /dev/null
+++ b/lightllm/common/basemodel/attention/nsa/dsv4_fp8_flashmla_sparse.py
@@ -0,0 +1,186 @@
+import dataclasses
+from typing import TYPE_CHECKING
+
+import torch
+
+from ..base_att import AttControl, BaseAttBackend, BaseDecodeAttState, BasePrefillAttState
+
+if TYPE_CHECKING:
+ from lightllm.common.basemodel.infer_struct import InferStateInfo
+
+
+# The current FlashMLA MODEL1 binary only instantiates these Q-head counts.
+_SUPPORTED_Q_HEADS = (64, 128)
+
+
+def get_dsv4_flashmla_padded_q_heads(q_head_num: int) -> int:
+ for supported_head_num in _SUPPORTED_Q_HEADS:
+ if q_head_num <= supported_head_num:
+ return supported_head_num
+ raise ValueError(f"FlashMLA does not support {q_head_num} local Q heads; supported counts: {_SUPPORTED_Q_HEADS}")
+
+
+def _view_cache(buffer: torch.Tensor, page_size: int) -> torch.Tensor:
+ from lightllm.common.kv_cache_mem_manager.deepseek4_mem_manager import DSV4_MLA_BYTES_PER_TOKEN
+
+ byte_num = page_size * DSV4_MLA_BYTES_PER_TOKEN
+ return buffer[:, :byte_num].view(buffer.shape[0], page_size, 1, DSV4_MLA_BYTES_PER_TOKEN)
+
+
+def _get_vllm_flashmla():
+ from vllm.v1.attention.ops import flashmla
+
+ return flashmla
+
+
+class DeepseekV4FlashMlaFp8SparseAttBackend(BaseAttBackend):
+ def __init__(self, model):
+ super().__init__(model=model)
+ self.real_q_head_num = model.config["num_attention_heads"] // model.tp_world_size_
+ self.padded_q_head_num = get_dsv4_flashmla_padded_q_heads(self.real_q_head_num)
+ self.compress_ratios = tuple(dict.fromkeys(model.config["compress_ratios"]))
+
+ def _flashmla_att(
+ self,
+ q: torch.Tensor,
+ packed_kv: torch.Tensor,
+ mem_manager,
+ nsa_dict: dict,
+ sched_meta,
+ flash_mla,
+ flashmla_out: torch.Tensor = None,
+ ) -> torch.Tensor:
+ from lightllm.common.kv_cache_mem_manager.deepseek4_mem_manager import (
+ DSV4_C128_PAGE_SIZE,
+ DSV4_C4_PAGE_SIZE,
+ DSV4_SWA_PAGE_SIZE,
+ )
+
+ ratio = nsa_dict["compress_ratio"]
+ extra_cache = None
+ if ratio == 4:
+ extra_page_size = DSV4_C4_PAGE_SIZE
+ elif ratio == 128:
+ extra_page_size = DSV4_C128_PAGE_SIZE
+ elif ratio != 0:
+ raise ValueError(f"unsupported DeepSeek-V4 compress ratio: {ratio}")
+ if ratio:
+ buffer = mem_manager.get_compressed_kv_buffer(nsa_dict["layer_index"])
+ extra_cache = _view_cache(buffer, extra_page_size)
+
+ kwargs = dict(
+ q=q.unsqueeze(1),
+ k_cache=_view_cache(packed_kv, DSV4_SWA_PAGE_SIZE),
+ block_table=None,
+ cache_seqlens=None,
+ head_dim_v=nsa_dict["head_dim_v"],
+ tile_scheduler_metadata=sched_meta,
+ num_splits=None,
+ softmax_scale=nsa_dict["softmax_scale"],
+ causal=False,
+ is_fp8_kvcache=True,
+ indices=nsa_dict["swa_indices"],
+ attn_sink=nsa_dict["attn_sink"],
+ topk_length=nsa_dict["swa_lengths"],
+ extra_k_cache=extra_cache,
+ extra_indices_in_kvcache=nsa_dict.get("extra_indices"),
+ extra_topk_length=nsa_dict.get("extra_lengths"),
+ )
+ if flashmla_out is not None:
+ kwargs["out"] = flashmla_out
+ full_out, _ = flash_mla.flash_mla_with_kvcache(**kwargs)
+ return full_out[:, 0, : self.real_q_head_num, :]
+
+ def create_att_prefill_state(self, infer_state: "InferStateInfo") -> "_PrefillAttState":
+ return _PrefillAttState(backend=self, infer_state=infer_state)
+
+ def create_att_decode_state(self, infer_state: "InferStateInfo") -> "_DecodeAttState":
+ return _DecodeAttState(backend=self, infer_state=infer_state)
+
+
+@dataclasses.dataclass
+class _PrefillAttState(BasePrefillAttState):
+ flashmla_sched_meta: dict = None
+ flash_mla: object = None
+
+ def init_state(self):
+ self.flash_mla = _get_vllm_flashmla()
+ self.flashmla_sched_meta = {}
+
+ def _get_sched_meta(self, compress_ratio: int):
+ if compress_ratio not in self.flashmla_sched_meta:
+ self.flashmla_sched_meta[compress_ratio] = self.flash_mla.get_mla_metadata()[0]
+ return self.flashmla_sched_meta[compress_ratio]
+
+ def prefill_att(
+ self,
+ q: torch.Tensor,
+ k: torch.Tensor,
+ v: torch.Tensor,
+ att_control: AttControl = AttControl(),
+ alloc_func=torch.empty,
+ *,
+ out: torch.Tensor = None,
+ ) -> torch.Tensor:
+ assert att_control.nsa_prefill, "nsa_prefill must be True for NSA prefill attention"
+ assert att_control.nsa_prefill_dict is not None, "nsa_prefill_dict is required"
+ nsa_dict = att_control.nsa_prefill_dict
+ if out is None:
+ out = alloc_func(
+ (q.shape[0], self.backend.real_q_head_num, nsa_dict["head_dim_v"]),
+ dtype=q.dtype,
+ device=q.device,
+ )
+ full_out = self.infer_state.dsv4_workspace.flashmla_prefill_full_out[: q.shape[0]]
+ out.copy_(
+ self.backend._flashmla_att(
+ q,
+ k,
+ self.infer_state.mem_manager,
+ nsa_dict,
+ self._get_sched_meta(nsa_dict["compress_ratio"]),
+ self.flash_mla,
+ flashmla_out=full_out,
+ )
+ )
+ return out
+
+
+@dataclasses.dataclass
+class _DecodeAttState(BaseDecodeAttState):
+ flashmla_sched_meta: dict = None
+ flash_mla: object = None
+
+ def init_state(self):
+ self.reset_sched_meta_for_capture()
+
+ def reset_sched_meta_for_capture(self):
+ # FlashMLA lazily binds extra-cache geometry, so ratios cannot share one sched-meta object.
+ self.flash_mla = _get_vllm_flashmla()
+ self.flashmla_sched_meta = {
+ ratio: self.flash_mla.get_mla_metadata()[0] for ratio in self.backend.compress_ratios
+ }
+
+ def decode_att(
+ self,
+ q: torch.Tensor,
+ k: torch.Tensor,
+ v: torch.Tensor,
+ att_control: AttControl = AttControl(),
+ alloc_func=torch.empty,
+ ) -> torch.Tensor:
+ assert att_control.nsa_decode, "nsa_decode must be True for NSA decode attention"
+ assert att_control.nsa_decode_dict is not None, "nsa_decode_dict is required"
+ nsa_dict = att_control.nsa_decode_dict
+ real_out = self.backend._flashmla_att(
+ q,
+ k,
+ self.infer_state.mem_manager,
+ nsa_dict,
+ self.flashmla_sched_meta[nsa_dict["compress_ratio"]],
+ self.flash_mla,
+ )
+ return real_out.contiguous()
+
+
+DSV4_NSA_BACKENDS = {"fp8kv_dsa": {"flashmla_sparse": DeepseekV4FlashMlaFp8SparseAttBackend}}
diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py
index 94f9d4c1a2..9f986fc873 100755
--- a/lightllm/common/basemodel/basemodel.py
+++ b/lightllm/common/basemodel/basemodel.py
@@ -42,6 +42,8 @@
class TpPartBaseModel:
+ is_mtp_draft_model = False
+
# weight class
pre_and_post_weight_class = None
transformer_weight_class = None
@@ -554,8 +556,10 @@ def _decode(
else:
infer_batch_size = model_input.batch_size
- if self.graph is not None and self.graph.can_run(
- batch_size=infer_batch_size, max_len_in_batch=model_input.max_kv_seq_len
+ if (
+ self.graph is not None
+ and not self.is_mtp_draft_model
+ and self.graph.can_run(batch_size=infer_batch_size, max_len_in_batch=model_input.max_kv_seq_len)
):
infer_batch_size = self.graph.find_closest_graph_batch_size(batch_size=infer_batch_size)
model_input = self._create_padded_decode_model_input(
@@ -600,6 +604,7 @@ def _decode(
def _context_forward(self, infer_state: InferStateInfo):
input_embs = self.pre_infer.context_forward(infer_state.input_ids, infer_state, self.pre_post_weight)
+ infer_state.mtp_draft_input_hiddens = None
if self.args.enable_dp_prefill_balance:
assert not self.args.enable_prefill_cudagraph, "not support now"
infer_state.prepare_prefill_dp_balance()
@@ -645,14 +650,20 @@ def prefill_func(input_tensors, infer_state):
last_input_embs = infer_state._all_to_all_unbalance_get(data=last_input_embs)
predict_logits = self.post_infer.token_forward(last_input_embs, infer_state, self.pre_post_weight)
+ mtp_main_output_hiddens = None
+ if isinstance(predict_logits, tuple):
+ predict_logits, mtp_main_output_hiddens = predict_logits
model_output = ModelOutput(logits=predict_logits)
# 特殊模型特殊模式的额外输出
if self.is_mtp_mode:
- input_embs = self.pre_infer._tpsp_allgather(input=input_embs, infer_state=infer_state)
- if infer_state.need_dp_prefill_balance:
- input_embs = infer_state._all_to_all_unbalance_get(data=input_embs)
- model_output.mtp_main_output_hiddens = input_embs.contiguous()
+ if mtp_main_output_hiddens is not None:
+ model_output.mtp_main_output_hiddens = mtp_main_output_hiddens.contiguous()
+ else:
+ input_embs = self.pre_infer._tpsp_allgather(input=input_embs, infer_state=infer_state)
+ if infer_state.need_dp_prefill_balance:
+ input_embs = infer_state._all_to_all_unbalance_get(data=input_embs)
+ model_output.mtp_main_output_hiddens = input_embs.contiguous()
# 在开启使用deepep的时候，需要调用clear_deepep_buffer做资源清理，没有启用的时候
# 该调用没有实际意义
@@ -664,6 +675,7 @@ def _token_forward(self, infer_state: InferStateInfo):
input_ids = infer_state.input_ids
cuda_input_ids = input_ids
input_embs = self.pre_infer.token_forward(cuda_input_ids, infer_state, self.pre_post_weight)
+ infer_state.mtp_draft_input_hiddens = None
input_embs = self.pre_infer._tpsp_sp_split(input=input_embs, infer_state=infer_state)
for i in range(self.layers_num):
@@ -671,16 +683,22 @@ def _token_forward(self, infer_state: InferStateInfo):
input_embs: torch.Tensor = layer.token_forward(input_embs, infer_state, self.trans_layers_weight[i])
last_input_embs = self.post_infer._tpsp_allgather(input=input_embs, infer_state=infer_state)
- predict_logits: torch.Tensor = self.post_infer.token_forward(
+ predict_logits = self.post_infer.token_forward(
last_input_embs, infer_state=infer_state, layer_weight=self.pre_post_weight
)
+ mtp_main_output_hiddens = None
+ if isinstance(predict_logits, tuple):
+ predict_logits, mtp_main_output_hiddens = predict_logits
model_output = ModelOutput(logits=predict_logits.contiguous())
# 特殊模型特殊模式的额外输出
if self.is_mtp_mode:
- input_embs = self.pre_infer._tpsp_allgather(input=input_embs, infer_state=infer_state)
- model_output.mtp_main_output_hiddens = input_embs.contiguous()
+ if mtp_main_output_hiddens is not None:
+ model_output.mtp_main_output_hiddens = mtp_main_output_hiddens.contiguous()
+ else:
+ input_embs = self.pre_infer._tpsp_allgather(input=input_embs, infer_state=infer_state)
+ model_output.mtp_main_output_hiddens = input_embs.contiguous()
# 在 cuda graph 模式下，输出需要转为 no ref tensor, 加强mem pool 的复用，降低显存的使用。
if infer_state.is_cuda_graph:
@@ -917,18 +935,31 @@ def _overlap_tpsp_context_forward(self, infer_state: InferStateInfo, infer_state
last_input_embs, last_input_embs1, infer_state, infer_state1, self.pre_post_weight
)
g_cache_manager.cache_env_out()
+ mtp_main_output_hiddens = None
+ mtp_main_output_hiddens1 = None
+ if isinstance(predict_logits, tuple):
+ predict_logits, mtp_main_output_hiddens = predict_logits
+ if isinstance(predict_logits1, tuple):
+ predict_logits1, mtp_main_output_hiddens1 = predict_logits1
model_output = ModelOutput(logits=predict_logits.contiguous())
model_output1 = ModelOutput(logits=predict_logits1.contiguous())
if self.is_mtp_mode:
- input_embs = self.pre_infer._tpsp_allgather(input=input_embs, infer_state=infer_state)
- input_embs1 = self.pre_infer._tpsp_allgather(input=input_embs1, infer_state=infer_state1)
- if infer_state.need_dp_prefill_balance:
- input_embs = infer_state._all_to_all_unbalance_get(data=input_embs)
- input_embs1 = infer_state1._all_to_all_unbalance_get(data=input_embs1)
- model_output.mtp_main_output_hiddens = input_embs.contiguous()
- model_output1.mtp_main_output_hiddens = input_embs1.contiguous()
+ if mtp_main_output_hiddens is not None:
+ model_output.mtp_main_output_hiddens = mtp_main_output_hiddens.contiguous()
+ else:
+ input_embs = self.pre_infer._tpsp_allgather(input=input_embs, infer_state=infer_state)
+ if infer_state.need_dp_prefill_balance:
+ input_embs = infer_state._all_to_all_unbalance_get(data=input_embs)
+ model_output.mtp_main_output_hiddens = input_embs.contiguous()
+ if mtp_main_output_hiddens1 is not None:
+ model_output1.mtp_main_output_hiddens = mtp_main_output_hiddens1.contiguous()
+ else:
+ input_embs1 = self.pre_infer._tpsp_allgather(input=input_embs1, infer_state=infer_state1)
+ if infer_state.need_dp_prefill_balance:
+ input_embs1 = infer_state1._all_to_all_unbalance_get(data=input_embs1)
+ model_output1.mtp_main_output_hiddens = input_embs1.contiguous()
return model_output, model_output1
@@ -955,15 +986,27 @@ def _overlap_tpsp_token_forward(self, infer_state: InferStateInfo, infer_state1:
predict_logits, predict_logits1 = self.post_infer.overlap_tpsp_token_forward(
last_input_embs, last_input_embs1, infer_state, infer_state1, self.pre_post_weight
)
+ mtp_main_output_hiddens = None
+ mtp_main_output_hiddens1 = None
+ if isinstance(predict_logits, tuple):
+ predict_logits, mtp_main_output_hiddens = predict_logits
+ if isinstance(predict_logits1, tuple):
+ predict_logits1, mtp_main_output_hiddens1 = predict_logits1
model_output = ModelOutput(logits=predict_logits.contiguous())
model_output1 = ModelOutput(logits=predict_logits1.contiguous())
if self.is_mtp_mode:
- input_embs = self.pre_infer._tpsp_allgather(input=input_embs, infer_state=infer_state)
- input_embs1 = self.pre_infer._tpsp_allgather(input=input_embs1, infer_state=infer_state1)
- model_output.mtp_main_output_hiddens = input_embs.contiguous()
- model_output1.mtp_main_output_hiddens = input_embs1.contiguous()
+ if mtp_main_output_hiddens is not None:
+ model_output.mtp_main_output_hiddens = mtp_main_output_hiddens.contiguous()
+ else:
+ input_embs = self.pre_infer._tpsp_allgather(input=input_embs, infer_state=infer_state)
+ model_output.mtp_main_output_hiddens = input_embs.contiguous()
+ if mtp_main_output_hiddens1 is not None:
+ model_output1.mtp_main_output_hiddens = mtp_main_output_hiddens1.contiguous()
+ else:
+ input_embs1 = self.pre_infer._tpsp_allgather(input=input_embs1, infer_state=infer_state1)
+ model_output1.mtp_main_output_hiddens = input_embs1.contiguous()
if infer_state.is_cuda_graph:
model_output.to_no_ref_tensor()
@@ -1171,13 +1214,7 @@ def _init_padded_req(self):
def _gen_special_model_input(self, token_num: int):
special_model_input = {}
- is_mtp_draft_model = (
- "Deepseek3MTPModel" in str(self.__class__)
- or "Qwen3MOEMTPModel" in str(self.__class__)
- or "MistralMTPModel" in str(self.__class__)
- or "Glm4MoeLiteMTPModel" in str(self.__class__)
- )
- if is_mtp_draft_model:
+ if self.is_mtp_draft_model:
special_model_input["mtp_draft_input_hiddens"] = torch.randn(
token_num, self.config["hidden_size"], dtype=self.data_type, device="cuda"
)
diff --git a/lightllm/common/basemodel/batch_objs.py b/lightllm/common/basemodel/batch_objs.py
index 1795ff9a82..9e56beeb4c 100644
--- a/lightllm/common/basemodel/batch_objs.py
+++ b/lightllm/common/basemodel/batch_objs.py
@@ -1,3 +1,4 @@
+import copy
import torch
from dataclasses import dataclass, field
from typing import Optional
@@ -42,6 +43,9 @@ class ModelInput:
multimodal_params: list = None
# cpu 变量
mem_indexes_cpu: torch.Tensor = None
+ b_req_idx_cpu: torch.Tensor = None
+ b_mtp_index_cpu: torch.Tensor = None
+ b_seq_len_cpu: torch.Tensor = None
# prefill 阶段使用的参数，但是不是推理过程使用的参数，是推理外部进行资源管理
# 的一些变量
b_prefill_has_output_cpu: List[bool] = None # 标记进行prefill的请求是否具有输出
@@ -52,6 +56,54 @@ class ModelInput:
# mtp_draft_input_hiddens 用于模型 mtp 模式下
# 的 draft 模型的输入
mtp_draft_input_hiddens: Optional[torch.Tensor] = None
+ # 主模型为 None: 准备所有 MTP 列；draft 首轮为 (): 无新槽；draft 追加后为 (k,): 只准备新槽。
+ mtp_decode_slot_prepare_indices: Optional[tuple] = None
+
+ def _capture_cpu_mirror(self, tensor_name: str, mirror_name: str):
+ tensor = getattr(self, tensor_name)
+ if tensor is not None and not tensor.is_cuda:
+ setattr(self, mirror_name, tensor)
+ return
+
+ def capture_cpu_mirrors(self):
+ self._capture_cpu_mirror("b_req_idx", "b_req_idx_cpu")
+ self._capture_cpu_mirror("b_mtp_index", "b_mtp_index_cpu")
+ self._capture_cpu_mirror("b_seq_len", "b_seq_len_cpu")
+ return
+
+ def make_mtp_draft_input(self):
+ model_input = copy.copy(self)
+ model_input.b_seq_len = self.b_seq_len.clone()
+ model_input.b_seq_len_cpu = self.b_seq_len_cpu.clone()
+ model_input.mtp_decode_slot_prepare_indices = ()
+ return model_input
+
+ def advance_mtp_decode_step(
+ self,
+ new_mem_indexes_cpu: torch.Tensor,
+ new_mem_indexes: torch.Tensor,
+ max_mtp_index: int,
+ ):
+ self.b_seq_len += 1
+ self.b_seq_len_cpu += 1
+ self.max_kv_seq_len += 1
+ self.mtp_decode_slot_prepare_indices = (max_mtp_index,)
+ slots_per_req = max_mtp_index + 1
+ self.mem_indexes_cpu = torch.cat(
+ [
+ self.mem_indexes_cpu.view(-1, slots_per_req)[:, 1:],
+ new_mem_indexes_cpu.view(-1, 1),
+ ],
+ dim=1,
+ ).view(-1)
+ self.mem_indexes = torch.cat(
+ [
+ self.mem_indexes.view(-1, slots_per_req)[:, 1:],
+ new_mem_indexes.view(-1, 1),
+ ],
+ dim=1,
+ ).view(-1)
+ return
def to_cuda(self):
if self.input_ids is not None:
@@ -82,6 +134,7 @@ def to_cuda(self):
self.b_shared_seq_len = self.b_shared_seq_len.cuda(non_blocking=True)
def __post_init__(self):
+ self.capture_cpu_mirrors()
self.check_input()
def check_input(self):
diff --git a/lightllm/common/basemodel/cuda_graph.py b/lightllm/common/basemodel/cuda_graph.py
index 782150661e..e1d96b744e 100644
--- a/lightllm/common/basemodel/cuda_graph.py
+++ b/lightllm/common/basemodel/cuda_graph.py
@@ -14,6 +14,20 @@
logger = init_logger(__name__)
+def _reset_att_state_sched_meta(infer_state: InferStateInfo):
+ # capture 前调用: warmup 趟用 copy.copy 浅拷贝共享 decode_att_state,其内部惰性初始化的
+ # 调度对象(如 FlashMLASchedMeta,首次内核调用时按当时数据规划并回写)会被 warmup 的
+ # dummy 负载锁定;若不重置,捕获趟将绑定为 dummy 规划的调度张量,所有 replay 都用错误
+ # 的 tile schedule(DSV4 实测 gsm8k 0.96 -> 0.74)。重置后规划发生在捕获区内,随 replay 重算。
+ for att_state in (infer_state.decode_att_state, infer_state.decode_att_state1):
+ if att_state is None:
+ continue
+ reset_fn = getattr(att_state, "reset_sched_meta_for_capture", None)
+ if reset_fn is not None:
+ reset_fn()
+ return
+
+
class CudaGraph:
# CudaGraph forward pass for the decoding stage.
@@ -94,6 +108,8 @@ def _capture_decode(self, decode_func, infer_state: InferStateInfo):
if param_name not in pure_para_set:
delattr(infer_state, param_name)
+ _reset_att_state_sched_meta(infer_state)
+
with torch.cuda.graph(graph_obj, pool=self.mempool):
model_output = decode_func(infer_state)
self.graph[batch_size] = (graph_obj, infer_state, model_output)
@@ -128,6 +144,9 @@ def _capture_decode_overlap(
if para_name not in pure_para_set1:
delattr(infer_state1, para_name)
+ _reset_att_state_sched_meta(infer_state)
+ _reset_att_state_sched_meta(infer_state1)
+
with torch.cuda.graph(graph_obj, pool=self.mempool):
model_output, model_output1 = decode_func(infer_state, infer_state1)
self.graph[batch_size] = (
diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py
index fca9b80fcf..bab48b9895 100644
--- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py
+++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py
@@ -68,12 +68,13 @@ def __init__(
auto_update_redundancy_expert=self.auto_update_redundancy_expert,
)
self.lock = threading.Lock()
+ self._moe_weight_finalized = False
self._create_weight()
def _init_config(self, network_config: Dict[str, Any]):
self.n_group = network_config.get("n_group", 0)
self.use_grouped_topk = self.n_group > 0
- self.norm_topk_prob = network_config["norm_topk_prob"]
+ self.norm_topk_prob = network_config.get("norm_topk_prob", False)
self.topk_group = network_config.get("topk_group", 0)
self.num_experts_per_tok = network_config["num_experts_per_tok"]
self.routed_scaling_factor = network_config.get("routed_scaling_factor", 1.0)
@@ -152,6 +153,26 @@ def experts(
per_expert_scale=self.per_expert_scale,
)
+ def experts_with_topk(
+ self,
+ input_tensor: torch.Tensor,
+ topk_weights: torch.Tensor,
+ topk_ids: torch.Tensor,
+ is_prefill: Optional[bool] = None,
+ clamp_limit: Optional[float] = None,
+ alloc_tensor_func=torch.empty,
+ ) -> torch.Tensor:
+ return self.fuse_moe_impl.fused_experts_with_topk(
+ input_tensor=input_tensor,
+ w13=self.w13,
+ w2=self.w2,
+ topk_weights=topk_weights,
+ topk_ids=topk_ids,
+ is_prefill=is_prefill,
+ clamp_limit=clamp_limit,
+ alloc_tensor_func=alloc_tensor_func,
+ )
+
def low_latency_dispatch(
self,
hidden_states: torch.Tensor,
@@ -280,7 +301,13 @@ def verify_load(self):
e_score_correction_bias_load_ok = (
True if self.e_score_correction_bias is None else getattr(self.e_score_correction_bias, "load_ok", False)
)
- return weight_load_ok and per_expert_scale_load_ok and e_score_correction_bias_load_ok
+ load_ok = weight_load_ok and per_expert_scale_load_ok and e_score_correction_bias_load_ok
+ if load_ok and not self._moe_weight_finalized:
+ finalize = getattr(self.quant_method, "finalize_moe_weight", None)
+ if finalize is not None:
+ finalize(self)
+ self._moe_weight_finalized = True
+ return load_ok
def _create_weight(self):
intermediate_size = self.split_inter_size
diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/__init__.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/__init__.py
index 67bb90e4ef..282c0abdce 100644
--- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/__init__.py
+++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/__init__.py
@@ -2,9 +2,15 @@
from .triton_impl import FuseMoeTriton
from .marlin_impl import FuseMoeMarlin
from .deepgemm_impl import FuseMoeDeepGEMM
+from .mxfp4_impl import FuseMoeMXFP4
def select_fuse_moe_impl(quant_method: QuantizationMethod, enable_ep_moe: bool):
+ if quant_method.method_name == "marlin-mxfp4w4a16-b32":
+ if enable_ep_moe:
+ raise RuntimeError("marlin-mxfp4w4a16-b32 does not support enable_ep_moe yet")
+ return FuseMoeMXFP4
+
if enable_ep_moe:
return FuseMoeDeepGEMM
diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py
index 4d4614c007..612c517155 100644
--- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py
+++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py
@@ -76,6 +76,8 @@ def _fused_experts(
topk_ids: torch.Tensor,
router_logits: Optional[torch.Tensor] = None,
is_prefill: Optional[bool] = None,
+ clamp_limit: Optional[float] = None,
+ alloc_tensor_func=torch.empty,
):
output = fused_experts(
hidden_states=input_tensor,
@@ -87,6 +89,8 @@ def _fused_experts(
quant_method=self.quant_method,
is_prefill=is_prefill,
previous_event=None, # for overlap
+ clamp_limit=clamp_limit,
+ alloc_tensor_func=alloc_tensor_func,
)
return output
@@ -192,6 +196,7 @@ def masked_group_gemm(
masked_m: torch.Tensor,
dtype: torch.dtype,
expected_m: int,
+ clamp_limit: Optional[float] = None,
):
w13_weight, w13_scale = w13.weight, w13.weight_scale
w2_weight, w2_scale = w2.weight, w2.weight_scale
@@ -204,6 +209,7 @@ def masked_group_gemm(
w2_weight,
w2_scale,
expected_m=expected_m,
+ clamp_limit=clamp_limit,
)
def prefilled_group_gemm(
diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/marlin_impl.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/marlin_impl.py
index 0094b09b1c..a30a669c18 100644
--- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/marlin_impl.py
+++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/marlin_impl.py
@@ -30,7 +30,10 @@ def _fused_experts(
topk_ids: torch.Tensor,
router_logits: Optional[torch.Tensor] = None,
is_prefill: Optional[bool] = None,
+ clamp_limit: Optional[float] = None,
+ alloc_tensor_func=torch.empty,
):
+ assert clamp_limit is None, "awq_marlin fused MoE does not support clamp_limit yet"
w1_weight, w1_scale, w1_zero_point = w13.weight, w13.weight_scale, w13.weight_zero_point
w2_weight, w2_scale, w2_zero_point = w2.weight, w2.weight_scale, w2.weight_zero_point
diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/mxfp4_impl.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/mxfp4_impl.py
new file mode 100644
index 0000000000..a7e19a9c80
--- /dev/null
+++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/mxfp4_impl.py
@@ -0,0 +1,45 @@
+import torch
+from typing import Optional
+
+from lightllm.common.quantization.quantize_method import WeightPack
+from .triton_impl import FuseMoeTriton
+
+
+class FuseMoeMXFP4(FuseMoeTriton):
+ def create_workspace(self):
+ return None
+
+ def _fused_experts(
+ self,
+ input_tensor: torch.Tensor,
+ w13: WeightPack,
+ w2: WeightPack,
+ topk_weights: torch.Tensor,
+ topk_ids: torch.Tensor,
+ router_logits: Optional[torch.Tensor] = None,
+ is_prefill: Optional[bool] = None,
+ clamp_limit: Optional[float] = None,
+ alloc_tensor_func=torch.empty,
+ ):
+ try:
+ from vllm.model_executor.layers.fused_moe.activation import MoEActivation
+ from vllm.model_executor.layers.fused_moe.experts.marlin_moe import fused_marlin_moe
+ from vllm.scalar_type import scalar_types
+ except Exception as e:
+ raise RuntimeError(f"MXFP4 fused MoE requires vLLM fused kernels, error={repr(e)}") from e
+
+ return fused_marlin_moe(
+ hidden_states=input_tensor.contiguous(),
+ w1=w13.weight,
+ w2=w2.weight,
+ bias1=None,
+ bias2=None,
+ w1_scale=w13.weight_scale,
+ w2_scale=w2.weight_scale,
+ topk_weights=topk_weights.to(torch.float32).contiguous(),
+ topk_ids=topk_ids.to(torch.long).contiguous(),
+ quant_type_id=scalar_types.float4_e2m1f.id,
+ global_num_experts=self.n_routed_experts,
+ activation=MoEActivation.SILU,
+ clamp_limit=clamp_limit,
+ )
diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py
index a0d30547a3..d8a3227236 100644
--- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py
+++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py
@@ -94,6 +94,8 @@ def _fused_experts(
topk_ids: torch.Tensor,
router_logits: Optional[torch.Tensor] = None,
is_prefill: bool = False,
+ clamp_limit: Optional[float] = None,
+ alloc_tensor_func=torch.empty,
):
w13_weight, w13_scale = w13.weight, w13.weight_scale
w2_weight, w2_scale = w2.weight, w2.weight_scale
@@ -111,9 +113,32 @@ def _fused_experts(
use_fp8_w8a8=use_fp8_w8a8,
w1_scale=w13_scale,
w2_scale=w2_scale,
+ limit=clamp_limit,
)
return input_tensor
+ def fused_experts_with_topk(
+ self,
+ input_tensor: torch.Tensor,
+ w13: WeightPack,
+ w2: WeightPack,
+ topk_weights: torch.Tensor,
+ topk_ids: torch.Tensor,
+ is_prefill: Optional[bool] = None,
+ clamp_limit: Optional[float] = None,
+ alloc_tensor_func=torch.empty,
+ ):
+ return self._fused_experts(
+ input_tensor=input_tensor,
+ w13=w13,
+ w2=w2,
+ topk_weights=topk_weights,
+ topk_ids=topk_ids,
+ is_prefill=is_prefill,
+ clamp_limit=clamp_limit,
+ alloc_tensor_func=alloc_tensor_func,
+ )
+
def __call__(
self,
input_tensor: torch.Tensor,
diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/mm_weight/mm_weight.py b/lightllm/common/basemodel/layer_weights/meta_weights/mm_weight/mm_weight.py
index 5021699143..1b966c5738 100644
--- a/lightllm/common/basemodel/layer_weights/meta_weights/mm_weight/mm_weight.py
+++ b/lightllm/common/basemodel/layer_weights/meta_weights/mm_weight/mm_weight.py
@@ -54,8 +54,23 @@ def __init__(
self.gen_weight_quant_param_names()
def mm(
- self, input_tensor: torch.Tensor, out: Optional[torch.Tensor] = None, use_custom_tensor_mananger: bool = True
+ self,
+ input_tensor: torch.Tensor,
+ out: Optional[torch.Tensor] = None,
+ use_custom_tensor_mananger: bool = True,
+ out_dtype: Optional[torch.dtype] = None,
) -> torch.Tensor:
+ if out_dtype is not None and not isinstance(self.quant_method, NoQuantization):
+ raise NotImplementedError(f"out_dtype is not supported for quant method {self.quant_method.method_name}")
+ if out_dtype is not None:
+ return self.quant_method.apply(
+ input_tensor,
+ self.mm_param,
+ out,
+ use_custom_tensor_mananger=use_custom_tensor_mananger,
+ bias=self.bias,
+ out_dtype=out_dtype,
+ )
return self.quant_method.apply(
input_tensor, self.mm_param, out, use_custom_tensor_mananger=use_custom_tensor_mananger, bias=self.bias
)
diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py
index cb2e370cb9..126e15326e 100644
--- a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py
+++ b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py
@@ -49,7 +49,7 @@ def check_ep_expert_dtype(quant_method: Any):
"EP MoE requires --expert_dtype to be one of ['fp8', 'fp4'], "
f"but the resolved fused_moe quant method is `{expert_dtype}`. "
"Please start with --expert_dtype fp8 or --expert_dtype fp4. "
- "Note that --expert_dtype fp4 is only supported on SM100 GPUs."
+ "Note that --expert_dtype fp4 with EP MoE is only supported on SM100 GPUs."
)
if expert_dtype == "deepgemm-fp4fp8-b32" and not is_sm100_gpu():
raise RuntimeError(
@@ -66,21 +66,25 @@ def masked_group_gemm(
w2: torch.Tensor,
w2_scale: torch.Tensor,
expected_m: int,
+ clamp_limit: Optional[float] = None,
+ alloc_tensor_func: Callable = torch.empty,
):
padded_m = recv_x[0].shape[1]
E, N, _ = w1.shape
block_size = 128
# groupgemm (masked layout)
- gemm_out_a = torch.empty((E, padded_m, N), device=recv_x[0].device, dtype=dtype)
+ gemm_out_a = alloc_tensor_func((E, padded_m, N), device=recv_x[0].device, dtype=dtype)
expected_m = min(expected_m, padded_m)
- qsilu_out_scale = torch.empty((E, padded_m, N // 2 // block_size), device=recv_x[0].device, dtype=torch.float32)
- qsilu_out = torch.empty((E, padded_m, N // 2), dtype=w1.dtype, device=recv_x[0].device)
+ qsilu_out_scale = alloc_tensor_func(
+ (E, padded_m, N // 2 // block_size), device=recv_x[0].device, dtype=torch.float32
+ )
+ qsilu_out = alloc_tensor_func((E, padded_m, N // 2), dtype=w1.dtype, device=recv_x[0].device)
# groupgemm (masked layout)
- gemm_out_b = torch.empty_like(recv_x[0], device=recv_x[0].device, dtype=dtype)
+ gemm_out_b = alloc_tensor_func(recv_x[0].shape, device=recv_x[0].device, dtype=dtype)
_deepgemm_grouped_fp8_nt_masked(recv_x, (w1, w1_scale), gemm_out_a, masked_m, expected_m)
- silu_and_mul_masked_post_quant_fwd(gemm_out_a, qsilu_out, qsilu_out_scale, block_size, masked_m)
+ silu_and_mul_masked_post_quant_fwd(gemm_out_a, qsilu_out, qsilu_out_scale, block_size, masked_m, limit=clamp_limit)
_deepgemm_grouped_fp8_nt_masked((qsilu_out, qsilu_out_scale), (w2, w2_scale), gemm_out_b, masked_m, expected_m)
return gemm_out_b
@@ -119,6 +123,8 @@ def mega_moe_impl(
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
quant_method: Any,
+ clamp_limit: Optional[float] = None,
+ alloc_tensor_func: Callable = torch.empty,
):
if not (HAS_DEEPGEMM and hasattr(deep_gemm, "fp8_fp4_mega_moe")):
raise RuntimeError("deep_gemm does not provide fp8-fp4 Mega MoE kernel")
@@ -149,13 +155,14 @@ def mega_moe_impl(
buffer.topk_idx[:num_tokens].copy_(topk_ids)
buffer.topk_weights[:num_tokens].copy_(topk_weights)
- output = torch.empty_like(hidden_states)
+ output = alloc_tensor_func(hidden_states.shape, device=hidden_states.device, dtype=hidden_states.dtype)
deep_gemm.fp8_fp4_mega_moe(
output,
l1_weights,
l2_weights,
buffer,
cumulative_local_expert_recv_stats=stats,
+ activation_clamp=clamp_limit,
)
return output
@@ -193,10 +200,21 @@ def fused_experts(
quant_method: Any,
is_prefill: Optional[bool],
previous_event: Optional[Any] = None,
+ clamp_limit: Optional[float] = None,
+ alloc_tensor_func: Callable = torch.empty,
):
check_ep_expert_dtype(quant_method)
if use_sm100_mega_moe(quant_method):
- return mega_moe_impl(hidden_states, w13, w2, topk_weights, topk_idx, quant_method)
+ return mega_moe_impl(
+ hidden_states,
+ w13,
+ w2,
+ topk_weights,
+ topk_idx,
+ quant_method,
+ clamp_limit=clamp_limit,
+ alloc_tensor_func=alloc_tensor_func,
+ )
buffer = dist_group_manager.ep_buffer if is_prefill else dist_group_manager.ep_low_latency_buffer
return fused_experts_impl(
@@ -214,6 +232,8 @@ def fused_experts(
w1_scale=w13.weight_scale,
w2_scale=w2.weight_scale,
previous_event=previous_event,
+ clamp_limit=clamp_limit,
+ alloc_tensor_func=alloc_tensor_func,
)
@@ -232,6 +252,8 @@ def fused_experts_impl(
w1_scale: Optional[torch.Tensor] = None,
w2_scale: Optional[torch.Tensor] = None,
previous_event: Optional[Any] = None,
+ clamp_limit: Optional[float] = None,
+ alloc_tensor_func: Callable = torch.empty,
):
# Check constraints.
assert hidden_states.shape[1] == w1.shape[2], "Hidden size mismatch"
@@ -256,7 +278,9 @@ def fused_experts_impl(
combined_x = None
if is_prefill:
- qinput_tensor, input_scale = per_token_group_quant_fp8(hidden_states, block_size_k, dtype=w1.dtype)
+ qinput_tensor, input_scale = per_token_group_quant_fp8(
+ hidden_states, block_size_k, dtype=w1.dtype, alloc_func=alloc_tensor_func
+ )
allocate_on_comm_stream = previous_event is not None
# normal dispatch
# recv_x [recive_num_tokens, hidden] recv_x_scale [recive_num_tokens, hidden // block_size]
@@ -298,7 +322,11 @@ def fused_experts_impl(
handle.num_recv_tokens_per_expert_list, dtype=torch.int32, pin_memory=True, device="cpu"
).cuda(non_blocking=True)
- expert_start_loc = torch.empty_like(num_recv_tokens_per_expert)
+ expert_start_loc = alloc_tensor_func(
+ num_recv_tokens_per_expert.shape,
+ device=num_recv_tokens_per_expert.device,
+ dtype=num_recv_tokens_per_expert.dtype,
+ )
ep_scatter(
recv_x[0],
@@ -311,23 +339,32 @@ def fused_experts_impl(
m_indices,
output_index,
)
+ recv_x = None
# groupgemm (contiguous layout)
- gemm_out_a = torch.empty((all_tokens, N), device=hidden_states.device, dtype=hidden_states.dtype)
+ gemm_storage = torch.empty(all_tokens * max(N, K), device=hidden_states.device, dtype=hidden_states.dtype)
+ gemm_out_a = gemm_storage[: all_tokens * N].view(all_tokens, N)
input_tensor[1] = tma_align_input_scale(input_tensor[1])
deepgemm_grouped_fp8_nt_contiguous(input_tensor, (w1, w1_scale), gemm_out_a, m_indices)
+ input_tensor = None
# silu_and_mul_fwd + qaunt
# TODO fused kernel
silu_out = torch.empty((all_tokens, N // 2), device=hidden_states.device, dtype=hidden_states.dtype)
- silu_and_mul_fwd(gemm_out_a.view(-1, N), silu_out)
+ silu_and_mul_fwd(gemm_out_a.view(-1, N), silu_out, limit=clamp_limit)
qsilu_out, qsilu_out_scale = per_token_group_quant_fp8(
- silu_out, block_size_k, dtype=w1.dtype, column_major_scales=True, scale_tma_aligned=True
+ silu_out,
+ block_size_k,
+ dtype=w1.dtype,
+ column_major_scales=True,
+ scale_tma_aligned=True,
)
+ gemm_out_a = None
+ silu_out = None
# groupgemm (contiguous layout)
- gemm_out_b = torch.empty((all_tokens, K), device=hidden_states.device, dtype=hidden_states.dtype)
+ gemm_out_b = gemm_storage[: all_tokens * K].view(all_tokens, K)
deepgemm_grouped_fp8_nt_contiguous((qsilu_out, qsilu_out_scale), (w2, w2_scale), gemm_out_b, m_indices)
@@ -365,7 +402,18 @@ def fused_experts_impl(
return_recv_hook=False,
)
# deepgemm
- gemm_out_b = masked_group_gemm(recv_x, masked_m, hidden_states.dtype, w1, w1_scale, w2, w2_scale, expected_m)
+ gemm_out_b = masked_group_gemm(
+ recv_x,
+ masked_m,
+ hidden_states.dtype,
+ w1,
+ w1_scale,
+ w2,
+ w2_scale,
+ expected_m,
+ clamp_limit=clamp_limit,
+ alloc_tensor_func=alloc_tensor_func,
+ )
# low latency combine
combined_x, event_overlap, hook = buffer.low_latency_combine(
gemm_out_b, topk_idx, topk_weights, handle, async_finish=False, return_recv_hook=False
diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul.py b/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul.py
index 45c7ea73c6..82fc9131c1 100644
--- a/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul.py
+++ b/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul.py
@@ -24,6 +24,7 @@ def _silu_and_mul_kernel_fast(
NEED_MASK: tl.constexpr,
layout: tl.constexpr = "blocked", # "blocked" or "interleaved"
USE_LIMIT_AND_ALPHA: tl.constexpr = False,
+ USE_LIMIT_ONLY: tl.constexpr = False,
USE_TANH_APPROXIMATE_GELU: tl.constexpr = False,
):
stride_input_m = tl.cast(stride_input_m, dtype=tl.int64)
@@ -76,6 +77,11 @@ def _silu_and_mul_kernel_fast(
mask=mask,
)
else:
+ if USE_LIMIT_ONLY:
+ # clamped swiglu (DeepSeek-V4 swiglu_limit): clamp 后接标准 silu，
+ # 无 gpt-oss 的 alpha 缩放与 (up+1)。
+ gate = tl.minimum(gate, limit)
+ up = tl.minimum(tl.maximum(up, -limit), limit)
if USE_TANH_APPROXIMATE_GELU:
# tanh-approx GELU, matching Gemma's gelu_pytorch_tanh MLP.
gate_cubed = gate * gate * gate
@@ -124,7 +130,8 @@ def silu_and_mul_fwd(
):
assert input.is_contiguous()
assert output.is_contiguous()
- assert (limit is None and alpha is None) or (limit is not None and alpha is not None)
+ # limit+alpha: gpt-oss 语义 (up+1)*silu(alpha*gate); 仅 limit: clamp 后标准 silu (DeepSeek-V4)
+ assert alpha is None or limit is not None
stride_input_m = input.stride(0)
stride_input_n = input.stride(1)
@@ -147,6 +154,7 @@ def silu_and_mul_fwd(
while triton.cdiv(size_m, BLOCK_M) > 8192:
BLOCK_M *= 2
USE_LIMIT_AND_ALPHA = limit is not None and alpha is not None
+ USE_LIMIT_ONLY = limit is not None and alpha is None
grid = (
triton.cdiv(size_n, BLOCK_N),
@@ -171,6 +179,7 @@ def silu_and_mul_fwd(
num_warps=num_warps,
layout=layout,
USE_LIMIT_AND_ALPHA=USE_LIMIT_AND_ALPHA,
+ USE_LIMIT_ONLY=USE_LIMIT_ONLY,
USE_TANH_APPROXIMATE_GELU=ffn_use_tanh_approximate_gelu(),
)
return
diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_mix_quant_ep.py b/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_mix_quant_ep.py
index aa91f15ed9..3c5c0de1de 100644
--- a/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_mix_quant_ep.py
+++ b/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_mix_quant_ep.py
@@ -24,8 +24,10 @@ def _silu_and_mul_post_quant_kernel(
size_n,
fp8_max,
fp8_min,
+ limit: tl.constexpr,
BLOCK_N: tl.constexpr,
NUM_STAGE: tl.constexpr,
+ USE_LIMIT_ONLY: tl.constexpr = False,
USE_TANH_APPROXIMATE_GELU: tl.constexpr = False,
):
expert_id = tl.program_id(2)
@@ -51,6 +53,9 @@ def _silu_and_mul_post_quant_kernel(
for token_index in tl.range(token_id, token_num_cur_expert, block_num_per_expert, num_stages=NUM_STAGE):
gate = tl.load(input_ptr_offs + token_index * stride_input_1, mask=offs_in_d < size_n, other=0.0).to(tl.float32)
up = tl.load(input_ptr_offs + token_index * stride_input_1 + size_n, mask=offs_in_d < size_n, other=0.0)
+ if USE_LIMIT_ONLY:
+ gate = tl.minimum(gate, limit)
+ up = tl.minimum(tl.maximum(up, -limit), limit)
if USE_TANH_APPROXIMATE_GELU:
gate_cubed = gate * gate * gate
tanh_arg = 0.7978845608028654 * (gate + 0.044715 * gate_cubed)
@@ -80,6 +85,7 @@ def silu_and_mul_masked_post_quant_fwd(
output_scale: torch.Tensor,
quant_group_size: int,
masked_m: torch.Tensor,
+ limit=None,
):
"""
input shape [expert_num, token_num_padded, hidden_dim]
@@ -133,8 +139,10 @@ def silu_and_mul_masked_post_quant_fwd(
size_n,
fp8_max,
fp8_min,
+ limit=limit,
BLOCK_N=BLOCK_N,
NUM_STAGE=NUM_STAGES,
+ USE_LIMIT_ONLY=limit is not None,
USE_TANH_APPROXIMATE_GELU=ffn_use_tanh_approximate_gelu(),
num_warps=num_warps,
)
diff --git a/lightllm/common/kv_cache_mem_manager/__init__.py b/lightllm/common/kv_cache_mem_manager/__init__.py
index 05544e149a..95f7e8ab76 100644
--- a/lightllm/common/kv_cache_mem_manager/__init__.py
+++ b/lightllm/common/kv_cache_mem_manager/__init__.py
@@ -4,6 +4,7 @@
from .ppl_int4kv_mem_manager import PPLINT4KVMemoryManager
from .deepseek2_mem_manager import Deepseek2MemoryManager
from .deepseek3_2mem_manager import Deepseek3_2MemoryManager
+from .deepseek4_mem_manager import DeepseekV4MemoryManager
from .fp8_per_token_group_quant_deepseek3_2mem_manager import FP8PerTokenGroupQuantDeepseek3_2MemoryManager
from .fp8_static_per_head_quant_mem_manager import FP8StaticPerHeadQuantMemManager
from .fp8_static_per_tensor_quant_mem_manager import FP8StaticPerTensorQuantMemManager
@@ -17,6 +18,7 @@
"PPLINT8KVMemoryManager",
"Deepseek2MemoryManager",
"Deepseek3_2MemoryManager",
+ "DeepseekV4MemoryManager",
"FP8PerTokenGroupQuantDeepseek3_2MemoryManager",
"FP8StaticPerHeadQuantMemManager",
"FP8StaticPerTensorQuantMemManager",
diff --git a/lightllm/common/kv_cache_mem_manager/allocator.py b/lightllm/common/kv_cache_mem_manager/allocator.py
index 850c158778..0179ed2714 100644
--- a/lightllm/common/kv_cache_mem_manager/allocator.py
+++ b/lightllm/common/kv_cache_mem_manager/allocator.py
@@ -3,13 +3,13 @@
from lightllm.utils.dist_utils import get_current_rank_in_node
from lightllm.utils.envs_utils import get_unique_server_name
from lightllm.utils.log_utils import init_logger
-from typing import Union, List
+from typing import Union, List, Optional
logger = init_logger(__name__)
class KvCacheAllocator:
- def __init__(self, size: int) -> None:
+ def __init__(self, size: int, shared_name: Optional[str] = None) -> None:
self.size = size
self.mem_state = torch.arange(
0, self.size, dtype=torch.int32, device="cpu", requires_grad=False, pin_memory=True
@@ -26,9 +26,11 @@ def __init__(self, size: int) -> None:
rank_in_node = get_current_rank_in_node()
# 用共享内存进行共享，router 模块读取进行精确的调度估计, nccl port 作为一个单机中单实列的标记。防止冲突。
- self.shared_can_use_token_num = SharedInt(
- f"{get_unique_server_name()}_mem_manger_can_use_token_num_{rank_in_node}"
- )
+ # shared_name 为 None 时使用主 kv 池的默认名(router 调度据此估算)；DeepSeek-V4 的压缩子池等
+ # 需要各自独立的计数器，传入区别于主池的唯一名，避免多个 allocator 写同一个共享计数器。
+ if shared_name is None:
+ shared_name = f"{get_unique_server_name()}_mem_manger_can_use_token_num_{rank_in_node}"
+ self.shared_can_use_token_num = SharedInt(shared_name)
self.shared_can_use_token_num.set_value(self.can_use_mem_size)
return
diff --git a/lightllm/common/kv_cache_mem_manager/deepseek4_mem_manager.py b/lightllm/common/kv_cache_mem_manager/deepseek4_mem_manager.py
new file mode 100644
index 0000000000..883bb59937
--- /dev/null
+++ b/lightllm/common/kv_cache_mem_manager/deepseek4_mem_manager.py
@@ -0,0 +1,781 @@
+import torch
+from typing import List, Optional, Union
+from .mem_manager import MemoryManager
+from .operator import DeepseekV4MemOperator
+from .allocator import KvCacheAllocator
+from lightllm.utils.dist_utils import get_current_rank_in_node
+from lightllm.utils.envs_utils import get_unique_server_name
+from lightllm.utils.log_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+# fp8_ds_mla packed-latent byte layout (ABI shared with the flash_mla extra-cache fork and
+# sglang/vllm): 448B NoPE fp8 + 64*2B RoPE bf16 + 7B ue8m0 scale + 1B pad = 584B per token,
+# stored in page slabs whose tail carries the per-token scale bytes.
+DSV4_MLA_NOPE_DIM = 448 # 448B
+DSV4_MLA_ROPE_DIM = 64 # 64 dim
+DSV4_MLA_HEAD_DIM = DSV4_MLA_NOPE_DIM + DSV4_MLA_ROPE_DIM # 512
+DSV4_MLA_QUANT_GROUP_SIZE = 64 # 64
+DSV4_MLA_SCALE_BYTES = DSV4_MLA_NOPE_DIM // DSV4_MLA_QUANT_GROUP_SIZE + 1 # 8 (7 ue8m0 + 1 pad)
+DSV4_MLA_BYTES_PER_TOKEN = DSV4_MLA_NOPE_DIM + DSV4_MLA_ROPE_DIM * 2 + DSV4_MLA_SCALE_BYTES # 584
+DSV4_MLA_DATA_BYTES_PER_TOKEN = DSV4_MLA_NOPE_DIM + DSV4_MLA_ROPE_DIM * 2 # 576
+DSV4_MLA_PAGE_ALIGN_BYTES = DSV4_MLA_DATA_BYTES_PER_TOKEN # 576
+DSV4_INDEXER_HEAD_DIM = 128 # 128
+DSV4_INDEXER_SCALE_BYTES = 4 # 4B fp32 scale
+DSV4_INDEXER_BYTES_PER_TOKEN = DSV4_INDEXER_HEAD_DIM + DSV4_INDEXER_SCALE_BYTES # 132
+DSV4_FP8_E4M3_MAX = 448.0 # 448.0
+DSV4_FP8_SCALE_MIN = 1e-4 # 1e-4
+DSV4_SWA_PAGE_SIZE = 128 # 128 slots/page
+DSV4_C4_PAGE_SIZE = 64 # 64 slots/page
+DSV4_C128_PAGE_SIZE = 2 # 2 slots/page
+DSV4_PROMPT_CACHE_PAGE_SIZE = DSV4_C4_PAGE_SIZE * 4 # 256 (= c4 ratio)
+# compressor state ring: c4 overlap 对为每页 2 个分组槽 × ratio 4 行;c128 离线聚合为每页 1 组。
+DSV4_C4_STATE_RING = 8 # 8 rows/page
+DSV4_C128_STATE_RING = 128 # 128 rows/page
+# swa 池占 full token 空间的比例(sglang DSV4 默认 swa_full_tokens_ratio=0.1 同值)。
+# 瞬时借页/驱逐走 swa 压力阀;池子大小仅按 ratio 切分,不再叠加结构性余量。
+DSV4_SWA_FULL_TOKENS_RATIO = 0.1 # 0.1
+
+
+def _ceil_div(a: int, b: int) -> int:
+ return (a + b - 1) // b
+
+
+class PackedPagePool:
+ """fp8_ds_mla 风格的 page-slab 存储: 每页前段连续放 token 的 data 字节，页尾放 per-token scale 字节。
+
+ 寻址是纯 token 槽位 (page = slot // page_size)，page 只是 scale-tail/对齐的物理打包技巧，
+ 不存在页粒度的分配。``write``/``read`` 是 torch 参考实现(单测 oracle)；生产写入走
+ triton packed writer(destindex_copy_kv_flashmla_dsv4 等)，kernel 直接消费 ``buffer``。
+ """
+
+ def __init__(
+ self,
+ size: int,
+ page_size: int,
+ layer_num: int,
+ data_bytes: int,
+ scale_bytes: int,
+ align_bytes: int = 1,
+ device: str = "cuda",
+ ):
+ self.size = size
+ self.page_size = page_size
+ self.layer_num = layer_num
+ self.data_bytes_per_token = data_bytes
+ self.scale_bytes_per_token = scale_bytes
+ self.bytes_per_token = data_bytes + scale_bytes
+ self.num_pages = _ceil_div(size + 1, page_size)
+ self.bytes_per_page = _ceil_div(page_size * self.bytes_per_token, align_bytes) * align_bytes
+ self.scale_offset_in_page = page_size * data_bytes
+ self.buffer = torch.zeros((layer_num, self.num_pages, self.bytes_per_page), dtype=torch.uint8, device=device)
+ self.HOLD_TOKEN_MEMINDEX = size
+
+ def get_layer_buffer(self, layer_index: int) -> torch.Tensor:
+ return self.buffer[layer_index]
+
+ def _loc_offsets(self, loc: torch.Tensor):
+ loc = loc.long()
+ page = torch.div(loc, self.page_size, rounding_mode="floor")
+ token = loc % self.page_size
+ page_base = page * self.bytes_per_page
+ data_offsets = page_base + token * self.data_bytes_per_token
+ scale_offsets = page_base + self.scale_offset_in_page + token * self.scale_bytes_per_token
+ return data_offsets, scale_offsets
+
+ def write(self, layer_index: int, loc: torch.Tensor, packed: torch.Tensor) -> None:
+ if loc.numel() == 0:
+ return
+ loc = loc.reshape(-1)
+ packed = packed.reshape(-1, self.bytes_per_token).contiguous()
+ flat = self.buffer[layer_index].view(-1)
+ data_offsets, scale_offsets = self._loc_offsets(loc)
+ data_range = torch.arange(self.data_bytes_per_token, device=loc.device)
+ scale_range = torch.arange(self.scale_bytes_per_token, device=loc.device)
+ flat[data_offsets.unsqueeze(1) + data_range.unsqueeze(0)] = packed[:, : self.data_bytes_per_token]
+ flat[scale_offsets.unsqueeze(1) + scale_range.unsqueeze(0)] = packed[:, self.data_bytes_per_token :]
+ return
+
+ def read(self, layer_index: int, loc: torch.Tensor) -> torch.Tensor:
+ loc = loc.reshape(-1)
+ if loc.numel() == 0:
+ return torch.empty((0, self.bytes_per_token), dtype=torch.uint8, device=self.buffer.device)
+ flat = self.buffer[layer_index].view(-1)
+ data_offsets, scale_offsets = self._loc_offsets(loc)
+ data_range = torch.arange(self.data_bytes_per_token, device=loc.device)
+ scale_range = torch.arange(self.scale_bytes_per_token, device=loc.device)
+ data = flat[data_offsets.unsqueeze(1) + data_range.unsqueeze(0)]
+ scale = flat[scale_offsets.unsqueeze(1) + scale_range.unsqueeze(0)]
+ return torch.cat([data, scale], dim=1).contiguous()
+
+
+class DeepseekV4MemoryManager(MemoryManager):
+ """DeepSeek-V4 KV cache: 窗口 latent(全层) + c4/c128 压缩 latent(压实层) + c4 indexer-K。
+
+ 与兄弟 manager 一致的 token-slot 设计；req 索引的表都在 DeepseekV4ReqManager。
+
+ - ``swa_pool``: 584B packed latent，所有层。池子小于 full token 空间；prep 阶段
+ ``alloc_swa_prefill/decode`` 按**页**(128 槽,位置对齐: slot(p)=page_base+p%128)分配，
+ 映射记录到 ``full_to_swa_indexs``(以 full token 槽位为键)。出窗槽位由 DeepseekV4ReqManager
+ 在 prep 阶段批量惰性回收(``evict_swa``,页存活计数减到 0 才整页归还)；full 槽位释放时
+ ``free`` 级联回收对应 swa 槽，所以 radix 驱逐/请求释放/暂停无需任何额外协议。
+ 页 allocator 触底时先走 swa free hook(radix 对 ref==0 节点 free)再 assert。
+ 没有 ring buffer，prefill chunk 大小不受 sliding_window 限制。
+ - ``c4_pool``/``c128_pool``: 压缩 latent，按 qwen3next 的层号压实手法只为压缩层建层；
+ c4 另带 packed indexer-K 池。槽位映射(``full_to_c4/c128_indexs``)以组末 token 的 full
+ 槽位为键(prep 阶段分配/scatter)，``free`` 级联回收，与 swa 完全同构。
+ - 写入走标准 operator 路径(``pack_mla_kv_to_cache``)，内部为 triton packed writer；
+ torch codecs 保留为 ABI 的可执行规格(单测 oracle)。
+ """
+
+ operator_class = DeepseekV4MemOperator
+
+ mla_nope_dim = DSV4_MLA_NOPE_DIM # 448
+ mla_rope_dim = DSV4_MLA_ROPE_DIM # 64
+ mla_head_dim = DSV4_MLA_HEAD_DIM # 512
+ mla_quant_group_size = DSV4_MLA_QUANT_GROUP_SIZE # 64
+ mla_scale_bytes = DSV4_MLA_SCALE_BYTES # 8
+ mla_bytes_per_token = DSV4_MLA_BYTES_PER_TOKEN # 584
+ indexer_head_dim_default = DSV4_INDEXER_HEAD_DIM # 128
+ indexer_bytes_per_token = DSV4_INDEXER_BYTES_PER_TOKEN # 132
+
+ def __init__(
+ self,
+ size,
+ dtype,
+ head_num,
+ head_dim,
+ layer_num,
+ compress_rates: List[int],
+ indexer_head_dim: int = 128,
+ max_request_num: Optional[int] = None,
+ sliding_window: Optional[int] = None,
+ swa_full_tokens_ratio: float = DSV4_SWA_FULL_TOKENS_RATIO,
+ always_copy=False,
+ mem_fraction=0.9,
+ ):
+ assert head_num == 1, "DeepSeek-V4 是 MLA(MQA)，dense latent 的 head_num 必须为 1"
+ assert head_dim == self.mla_head_dim, f"DeepSeek-V4 packed KV 期望 head_dim={self.mla_head_dim}"
+ assert (
+ indexer_head_dim == self.indexer_head_dim_default
+ ), f"DeepSeek-V4 packed indexer-K 期望 indexer_head_dim={self.indexer_head_dim_default}"
+ assert len(compress_rates) == layer_num, f"compress_rates 长度 {len(compress_rates)} 必须等于 layer_num {layer_num}"
+ assert all(r in (0, 4, 128) for r in compress_rates), "compress_rates 取值只能是 0/4/128"
+
+ self.compress_rates = list(compress_rates)
+ self.n_c4 = sum(1 for r in self.compress_rates if r == 4)
+ self.n_c128 = sum(1 for r in self.compress_rates if r == 128)
+ self.indexer_head_dim = indexer_head_dim
+ self.max_request_num = max_request_num
+ self.sliding_window = sliding_window
+ self.swa_full_tokens_ratio = float(swa_full_tokens_ratio)
+
+ # 全局层号 -> 各压缩池内的压实层号(同 qwen3next 的层号压实手法)
+ self.layer_to_c4_idx = {}
+ self.layer_to_c128_idx = {}
+ c4 = c128 = 0
+ for lid, r in enumerate(self.compress_rates):
+ if r == 4:
+ self.layer_to_c4_idx[lid] = c4
+ c4 += 1
+ elif r == 128:
+ self.layer_to_c128_idx[lid] = c128
+ c128 += 1
+
+ super().__init__(size, dtype, head_num, head_dim, layer_num, always_copy, mem_fraction)
+
+ # ------------------------------------------------------------------ sizing
+ def _planned_swa_size(self, full_size: int) -> int:
+ return _ceil_div(int(full_size * self.swa_full_tokens_ratio), DSV4_SWA_PAGE_SIZE) * DSV4_SWA_PAGE_SIZE
+
+ @staticmethod
+ def _paged_state_rows(num_swa_pages: int, ring: int, ratio: int) -> int:
+ rows = num_swa_pages * ring + ring + 1
+ return _ceil_div(rows, ratio) * ratio
+
+ @staticmethod
+ def _init_state_sentinel(buffer: torch.Tensor) -> None:
+ half = buffer.shape[-1] // 2
+ buffer[:, -1, :half].zero_()
+ buffer[:, -1, half:].fill_(float("-inf"))
+ return
+
+ def get_cell_size(self):
+ kv_bytes = self.mla_bytes_per_token
+ indexer_bytes = self.indexer_bytes_per_token
+ state_dtype_bytes = torch._utils._element_size(torch.float32)
+ c4_state_width = 4 * self.head_dim + 4 * self.indexer_head_dim
+ c128_state_width = 2 * self.head_dim
+ c4_state_bytes = DSV4_C4_STATE_RING / DSV4_SWA_PAGE_SIZE * c4_state_width * state_dtype_bytes * self.n_c4
+ c128_state_bytes = (
+ DSV4_C128_STATE_RING / DSV4_SWA_PAGE_SIZE * c128_state_width * state_dtype_bytes * self.n_c128
+ )
+ swa_slot = kv_bytes * self.layer_num + c4_state_bytes + c128_state_bytes
+ compressed = (kv_bytes + indexer_bytes) * self.n_c4 / 4 + kv_bytes * self.n_c128 / 128
+
+ return swa_slot * self.swa_full_tokens_ratio + compressed
+
+ # ------------------------------------------------------------------ buffers
+ def _init_buffers(self, size, dtype, head_num, head_dim, layer_num):
+ rank_in_node = get_current_rank_in_node()
+ server = get_unique_server_name()
+
+ self.swa_size = self._planned_swa_size(size)
+ self.swa_pool = PackedPagePool(
+ size=self.swa_size,
+ page_size=DSV4_SWA_PAGE_SIZE,
+ layer_num=layer_num,
+ data_bytes=DSV4_MLA_DATA_BYTES_PER_TOKEN,
+ scale_bytes=self.mla_scale_bytes,
+ align_bytes=DSV4_MLA_PAGE_ALIGN_BYTES,
+ )
+ # 注意: 该别名是 page 索引([layer, num_pages, bytes_per_page])而非 token 索引，
+ # 只允许 get_att_input_params 的消费者使用；token 索引语义的继承接口已显式 fence。
+ self.kv_buffer = self.swa_pool.buffer
+ # 页粒度分配(页 = 128 槽,位置对齐): 槽位不变式 slot(p) = page_base + p%128。
+ # swa_size 整页对齐 ⇒ HOLD 槽(swa_size)独占池子最后一个物理页,永不参与分配。
+ self.swa_num_pages = self.swa_size // DSV4_SWA_PAGE_SIZE
+ self.swa_page_allocator = KvCacheAllocator(
+ self.swa_num_pages, shared_name=f"{server}_dsv4_swa_can_use_page_num_{rank_in_node}"
+ )
+ # 页存活计数 = 指向该页的有效 full_to_swa 行数;减到 0 归还 allocator(出窗逐 token
+ # 回收下,「部分出窗页」计数 > 0 自然受保护)。下标含 HOLD 页(只读不增减)。
+ self.swa_page_live_count = torch.zeros((self.swa_pool.num_pages,), dtype=torch.int32, device="cuda")
+ # swa free hook(可选): 页 allocator 触底时回调(radix 对 ref==0 节点 free swa 页),
+ # 由 backend 在 radix cache 创建后 register;assert 仍是最后防线。
+ self._free_radix_unreferenced_swa_fn = None
+ self.full_to_swa_indexs = torch.full((size + 1,), -1, dtype=torch.int32, device="cuda")
+ self.full_to_swa_indexs[size] = self.swa_pool.HOLD_TOKEN_MEMINDEX
+
+ self.c4_size = _ceil_div(size, 4)
+ self.c128_size = _ceil_div(size, 128)
+ self.c4_pool: Optional[PackedPagePool] = None
+ self.c4_indexer_pool: Optional[PackedPagePool] = None
+ self.c4_allocator: Optional[KvCacheAllocator] = None
+ self.c4_page_allocator: Optional[KvCacheAllocator] = None
+ self.c4_page_live_count: Optional[torch.Tensor] = None
+ self.c128_pool: Optional[PackedPagePool] = None
+ self.c128_allocator: Optional[KvCacheAllocator] = None
+ self.c4_state_buffer: Optional[torch.Tensor] = None
+ self.c4_indexer_state_buffer: Optional[torch.Tensor] = None
+ self.c128_state_buffer: Optional[torch.Tensor] = None
+ # 压缩槽映射: 键 = 组末 token(位置 (g+1)%ratio==0)的 full 槽位,值 = 压缩池槽位。
+ # 与 full_to_swa_indexs 同构: radix 持有 full 槽 => 映射行存活,free 级联回收。
+ self.full_to_c4_indexs: Optional[torch.Tensor] = None
+ self.full_to_c128_indexs: Optional[torch.Tensor] = None
+ if self.n_c4 > 0:
+ self.c4_pool = PackedPagePool(
+ size=self.c4_size,
+ page_size=DSV4_C4_PAGE_SIZE,
+ layer_num=self.n_c4,
+ data_bytes=DSV4_MLA_DATA_BYTES_PER_TOKEN,
+ scale_bytes=self.mla_scale_bytes,
+ align_bytes=DSV4_MLA_PAGE_ALIGN_BYTES,
+ )
+ self.c4_indexer_pool = PackedPagePool(
+ size=self.c4_size,
+ page_size=DSV4_C4_PAGE_SIZE,
+ layer_num=self.n_c4,
+ data_bytes=self.indexer_head_dim,
+ scale_bytes=DSV4_INDEXER_SCALE_BYTES,
+ )
+ self.c4_num_pages = self.c4_size // DSV4_C4_PAGE_SIZE
+ assert self.c4_num_pages > 0, "DeepSeek-V4 c4 pool must have at least one usable full page"
+ self.c4_page_allocator = KvCacheAllocator(
+ self.c4_num_pages, shared_name=f"{server}_dsv4_c4_can_use_page_num_{rank_in_node}"
+ )
+ self.c4_page_live_count = torch.zeros((self.c4_pool.num_pages,), dtype=torch.int32, device="cuda")
+ self.full_to_c4_indexs = torch.full((size + 1,), -1, dtype=torch.int32, device="cuda")
+ self.full_to_c4_indexs[size] = self.c4_pool.HOLD_TOKEN_MEMINDEX
+ # c4 compressor 在途状态(attention + indexer): swa 页派生寻址(翻译③),随 swa 页
+ # 生灭 -> radix 命中零拷贝续算。行数 = 页数*ring + ring(HOLD 页) + 1(哨兵),
+ # 取整到 ratio;末行哨兵 kv=0/score=-inf(KVAndScore.clear 语义),其余行由内核在
+ # 组起点覆写,无需按页清零。last_dim = 2*coff*head_dim(overlap coff=2)。
+ state_rows = self._paged_state_rows(self.swa_num_pages, DSV4_C4_STATE_RING, 4)
+ self.c4_state_buffer = torch.zeros(
+ (self.n_c4, state_rows, 4 * self.head_dim), dtype=torch.float32, device="cuda"
+ )
+ self.c4_indexer_state_buffer = torch.zeros(
+ (self.n_c4, state_rows, 4 * self.indexer_head_dim), dtype=torch.float32, device="cuda"
+ )
+ for buf in (self.c4_state_buffer, self.c4_indexer_state_buffer):
+ self._init_state_sentinel(buf)
+ if self.n_c128 > 0:
+ self.c128_pool = PackedPagePool(
+ size=self.c128_size,
+ page_size=DSV4_C128_PAGE_SIZE,
+ layer_num=self.n_c128,
+ data_bytes=DSV4_MLA_DATA_BYTES_PER_TOKEN,
+ scale_bytes=self.mla_scale_bytes,
+ align_bytes=DSV4_MLA_PAGE_ALIGN_BYTES,
+ )
+ self.c128_allocator = KvCacheAllocator(
+ self.c128_size, shared_name=f"{server}_dsv4_c128_can_use_token_num_{rank_in_node}"
+ )
+ self.full_to_c128_indexs = torch.full((size + 1,), -1, dtype=torch.int32, device="cuda")
+ self.full_to_c128_indexs[size] = self.c128_pool.HOLD_TOKEN_MEMINDEX
+ # c128 compressor 在途状态: 与 c4 同样由 full->swa 推导行号,但 ring=128 且无 overlap。
+ # last_dim = 2*head_dim;末行是 swa 缺失/出窗时读取的哨兵。
+ state_rows = self._paged_state_rows(self.swa_num_pages, DSV4_C128_STATE_RING, 128)
+ self.c128_state_buffer = torch.zeros(
+ (self.n_c128, state_rows, 2 * self.head_dim), dtype=torch.float32, device="cuda"
+ )
+ self._init_state_sentinel(self.c128_state_buffer)
+
+ logger.info(
+ f"DeepseekV4MemoryManager pools: full_tokens={size} swa={self.swa_size}({self.swa_num_pages}p) "
+ f"c4={self.c4_size}(L={self.n_c4}) c128={self.c128_size}(L={self.n_c128}) "
+ f"packed_kv_bytes={self.mla_bytes_per_token} indexer_bytes={self.indexer_bytes_per_token}"
+ )
+
+ # ------------------------------------------------------------------ buffer accessors
+ def get_att_input_params(self, layer_index: int):
+ return self.swa_pool.get_layer_buffer(layer_index)
+
+ def _pool_and_local_layer(self, layer_index: int):
+ r = self.compress_rates[layer_index]
+ if r == 4:
+ return self.c4_pool, self.layer_to_c4_idx[layer_index]
+ if r == 128:
+ return self.c128_pool, self.layer_to_c128_idx[layer_index]
+ raise AssertionError(f"layer {layer_index} (rate {r}) 不是压缩层，没有压缩池")
+
+ def get_compressed_kv_buffer(self, layer_index: int) -> torch.Tensor:
+ pool, local_layer = self._pool_and_local_layer(layer_index)
+ return pool.get_layer_buffer(local_layer)
+
+ def get_indexer_k_buffer(self, layer_index: int) -> torch.Tensor:
+ assert self.compress_rates[layer_index] == 4, "只有 c4(CSA) 层有 indexer-K"
+ return self.c4_indexer_pool.get_layer_buffer(self.layer_to_c4_idx[layer_index])
+
+ def get_c4_state_buffer(self, layer_index: int) -> torch.Tensor:
+ assert self.compress_rates[layer_index] == 4, "只有 c4(CSA) 层有 paged compressor state"
+ return self.c4_state_buffer[self.layer_to_c4_idx[layer_index]]
+
+ def get_c4_indexer_state_buffer(self, layer_index: int) -> torch.Tensor:
+ assert self.compress_rates[layer_index] == 4, "只有 c4(CSA) 层有 paged indexer state"
+ return self.c4_indexer_state_buffer[self.layer_to_c4_idx[layer_index]]
+
+ def get_c128_state_buffer(self, layer_index: int) -> torch.Tensor:
+ assert self.compress_rates[layer_index] == 128, "只有 c128(HCA) 层有 paged compressor state"
+ return self.c128_state_buffer[self.layer_to_c128_idx[layer_index]]
+
+ # ------------------------------------------------------------------ swa slot lifecycle
+ def register_swa_free_hook(self, fn) -> None:
+ """fn(need_pages): 在页 allocator 不足时尝试腾页(radix 对 ref==0 节点 free swa)。"""
+ self._free_radix_unreferenced_swa_fn = fn
+ return
+
+ def _alloc_swa_pages(self, need_pages: int) -> torch.Tensor:
+ if need_pages > self.swa_page_allocator.can_use_mem_size and self._free_radix_unreferenced_swa_fn is not None:
+ self._free_radix_unreferenced_swa_fn(need_pages - self.swa_page_allocator.can_use_mem_size)
+ return self.swa_page_allocator.alloc(need_pages)
+
+ def _update_swa_page_counts(self, swa_slots: torch.Tensor, delta: int) -> torch.Tensor:
+ """按 slot 所在页更新存活计数，返回逐 slot 的页号。"""
+ pages = torch.div(swa_slots, DSV4_SWA_PAGE_SIZE, rounding_mode="floor")
+ ones = torch.full(pages.shape, delta, dtype=torch.int32, device=pages.device)
+ self.swa_page_live_count.index_add_(0, pages, ones)
+ return pages
+
+ def alloc_swa_prefill(
+ self,
+ mem_indexes: torch.Tensor,
+ req_to_token_indexs: torch.Tensor,
+ req_list: List[int],
+ ready_list: List[int],
+ seq_list: List[int],
+ ) -> None:
+ """prefill prep: 为各请求位置 [ready, seq) 的新 token 分配位置对齐的 swa 槽。
+
+ 槽位不变式: slot(p) = page_base(p 所在页) + p%128,page_base % 128 == 0。
+ 续页(start 非整页,只可能是首页)的 base 从上一 token 的映射派生
+ (full_to_swa[req_to_token[req, start-1]],该 token 必在保留窗内);其余页全新分配。
+ radix 命中(ready 必 128 对齐)的借用方从全新页开始,与节点持有页天然不相交。
+ 当前 chunk 的 full 槽直接来自 generic preprocess 分配的 mem_indexes，因此不依赖
+ req_to_token_indexs 已完成当前 chunk 的 scatter；只有续页的上一 token 查询旧 req 行。
+ """
+ page = DSV4_SWA_PAGE_SIZE
+ hold_req_id = self.max_request_num # padding 行的请求 id(req_manager.HOLD_REQUEST_ID)
+
+ segs = [] # (req_idx, start, end, mem_offset, n_new_pages, has_cont_page)
+ total_new_pages = 0
+ mem_offset = 0
+ for req_idx, start, end in zip(req_list, ready_list, seq_list):
+ q_len = end - start
+ if req_idx == hold_req_id or end <= start:
+ mem_offset += q_len
+ continue
+ first_new_page = _ceil_div(start, page)
+ n_new = max(0, (end - 1) // page - first_new_page + 1)
+ segs.append((req_idx, start, end, mem_offset, n_new, start % page != 0))
+ total_new_pages += n_new
+ mem_offset += q_len
+ if not segs:
+ return
+
+ device = self.full_to_swa_indexs.device
+ mem_indexes = mem_indexes.reshape(-1)
+ new_pages = self._alloc_swa_pages(total_new_pages).to(device, non_blocking=True) if total_new_pages else None
+ page_cursor = 0
+ for req_idx, start, end, mem_start, n_new, has_cont in segs:
+ positions = torch.arange(start, end, dtype=torch.int32, device=device)
+ page_local = torch.div(positions, page, rounding_mode="floor") - start // page
+ bases = torch.empty(((end - 1) // page - start // page + 1,), dtype=torch.int32, device=device)
+ if has_cont:
+ prev_slot = self.full_to_swa_indexs[req_to_token_indexs[req_idx, start - 1]]
+ bases[0] = prev_slot - (start - 1) % page
+ if n_new:
+ bases[1 if has_cont else 0 :] = new_pages[page_cursor : page_cursor + n_new] * page
+ page_cursor += n_new
+ slots = bases[page_local] + positions % page
+ full_slots = mem_indexes[mem_start : mem_start + end - start]
+ self.full_to_swa_indexs[full_slots] = slots
+ self._update_swa_page_counts(slots, 1)
+ return
+
+ def alloc_swa_decode(
+ self,
+ req_list: List[int],
+ seq_list: List[int],
+ mem_indexes: torch.Tensor,
+ prev_full_indexes: torch.Tensor,
+ ) -> None:
+ """decode prep: 本步 token(位置 seq-1)的 swa 槽。整页起点开新页,否则上一 token 槽 +1
+ (位置对齐不变式保证同页连续)。scatter 目标用当前步 mem_indexes。
+
+ 调用方传入每行前一 token 的 full 槽；MTP step>0 可直接使用同批前一列。"""
+ page = DSV4_SWA_PAGE_SIZE
+ hold_req_id = self.max_request_num
+ cont_rows, new_rows = [], []
+ for i, (req_idx, seq_len) in enumerate(zip(req_list, seq_list)):
+ if req_idx == hold_req_id or seq_len <= 0:
+ continue
+ if (seq_len - 1) % page == 0:
+ new_rows.append(i)
+ else:
+ cont_rows.append(i)
+ mem_indexes = mem_indexes.reshape(-1)
+ if cont_rows:
+ prev_full = prev_full_indexes.reshape(-1)[cont_rows]
+ prev_slots = self.full_to_swa_indexs[prev_full]
+ slots = prev_slots + 1
+ self.full_to_swa_indexs[mem_indexes[cont_rows]] = slots
+ self._update_swa_page_counts(slots, 1)
+ if new_rows:
+ pages = self._alloc_swa_pages(len(new_rows)).to(self.full_to_swa_indexs.device, non_blocking=True)
+ slots = pages * page
+ self.full_to_swa_indexs[mem_indexes[new_rows]] = slots
+ self._update_swa_page_counts(slots, 1)
+ return
+
+ def evict_swa(self, full_slots: torch.Tensor) -> None:
+ """回收 full 槽位对应的 swa 槽(出窗惰性回收 / free 级联 / 压力阀共用)。
+ 未映射(-1)的槽位跳过;页计数减到 0 时整页归还 allocator。"""
+ if full_slots.numel() == 0:
+ return
+ full_slots = full_slots.to(self.full_to_swa_indexs.device, non_blocking=True).reshape(-1)
+ full_slots = torch.unique(full_slots[full_slots != self.HOLD_TOKEN_MEMINDEX])
+ if full_slots.numel() == 0:
+ return
+ swa_slots = self.full_to_swa_indexs[full_slots]
+ valid = swa_slots >= 0
+ valid_slots = swa_slots[valid]
+ if valid_slots.numel() == 0:
+ return
+ self.full_to_swa_indexs[full_slots[valid]] = -1
+ touched = torch.unique(self._update_swa_page_counts(valid_slots, -1))
+ empty = touched[self.swa_page_live_count[touched] == 0]
+ if empty.numel() > 0:
+ self.swa_page_allocator.free(empty.to(torch.int32))
+ return
+
+ def _evict_compress(self, full_slots: torch.Tensor, mapping: torch.Tensor, allocator: KvCacheAllocator) -> None:
+ full_slots = full_slots.to(mapping.device, non_blocking=True).reshape(-1)
+ # 去重: 同批重复槽会 gather 出重复的压缩槽 -> allocator 双重释放(free 已去重,直呼叫方防御)。
+ full_slots = torch.unique(full_slots[full_slots != self.HOLD_TOKEN_MEMINDEX])
+ if full_slots.numel() == 0:
+ return
+ slots = mapping[full_slots]
+ valid = slots >= 0
+ valid_slots = slots[valid]
+ if valid_slots.numel() == 0:
+ return
+ allocator.free(valid_slots)
+ mapping[full_slots[valid]] = -1
+ return
+
+ def alloc_c4_pages(self, need_pages: int) -> torch.Tensor:
+ assert self.c4_page_allocator is not None, "DeepSeek-V4 c4 page allocator is not initialized"
+ return self.c4_page_allocator.alloc(need_pages)
+
+ def count_c4_slots(self, c4_slots: torch.Tensor, delta: int) -> torch.Tensor:
+ """按 c4 slot 所在页更新存活计数，返回逐 slot 的页号。"""
+ assert self.c4_page_live_count is not None, "DeepSeek-V4 c4 page live count is not initialized"
+ pages = torch.div(c4_slots, DSV4_C4_PAGE_SIZE, rounding_mode="floor")
+ ones = torch.full(pages.shape, delta, dtype=torch.int32, device=pages.device)
+ self.c4_page_live_count.index_add_(0, pages, ones)
+ return pages
+
+ def evict_c4(self, full_slots: torch.Tensor) -> None:
+ """回收 full 槽位(组末 token)映射的 c4 槽。非组末/未映射(-1)的槽位跳过。"""
+ if self.c4_page_allocator is None or full_slots.numel() == 0:
+ return
+ full_slots = full_slots.to(self.full_to_c4_indexs.device, non_blocking=True).reshape(-1)
+ full_slots = torch.unique(full_slots[full_slots != self.HOLD_TOKEN_MEMINDEX])
+ if full_slots.numel() == 0:
+ return
+ slots = self.full_to_c4_indexs[full_slots]
+ valid = slots >= 0
+ valid_slots = slots[valid]
+ if valid_slots.numel() == 0:
+ return
+ self.full_to_c4_indexs[full_slots[valid]] = -1
+ touched = torch.unique(self.count_c4_slots(valid_slots, -1))
+ empty = touched[self.c4_page_live_count[touched] == 0]
+ if empty.numel() > 0:
+ self.c4_page_allocator.free(empty.to(torch.int32))
+ return
+
+ def evict_c128(self, full_slots: torch.Tensor) -> None:
+ """回收 full 槽位(组末 token)映射的 c128 槽。非组末/未映射(-1)的槽位跳过。"""
+ if self.c128_allocator is None or full_slots.numel() == 0:
+ return
+ self._evict_compress(full_slots, self.full_to_c128_indexs, self.c128_allocator)
+ return
+
+ # ------------------------------------------------------------------ alloc/free (cascade)
+ def free(self, free_index: Union[torch.Tensor, List[int]]) -> None:
+ """释放 full token 槽位，级联回收其 swa 槽与 c4/c128 压缩槽。radix 驱逐、请求释放/暂停都走这里。
+
+ 先对 full 槽去重: 同批重复槽位会让映射 gather 出重复的压缩/swa 槽，导致 allocator 双重释放。"""
+ if isinstance(free_index, list):
+ free_index = torch.tensor(free_index, dtype=torch.int64)
+ if free_index.numel() > 0:
+ free_index = torch.unique(free_index)
+ self.evict_swa(free_index)
+ self.evict_c4(free_index)
+ self.evict_c128(free_index)
+ super().free(free_index)
+ return
+
+ def free_all(self):
+ super().free_all()
+ self.swa_page_allocator.free_all()
+ self.swa_page_live_count.zero_()
+ self.full_to_swa_indexs.fill_(-1)
+ self.full_to_swa_indexs[self.HOLD_TOKEN_MEMINDEX] = self.swa_pool.HOLD_TOKEN_MEMINDEX
+ if self.c4_page_allocator is not None:
+ self.c4_page_allocator.free_all()
+ self.c4_page_live_count.zero_()
+ self.full_to_c4_indexs.fill_(-1)
+ self.full_to_c4_indexs[self.HOLD_TOKEN_MEMINDEX] = self.c4_pool.HOLD_TOKEN_MEMINDEX
+ if self.c128_allocator is not None:
+ self.c128_allocator.free_all()
+ self.full_to_c128_indexs.fill_(-1)
+ self.full_to_c128_indexs[self.HOLD_TOKEN_MEMINDEX] = self.c128_pool.HOLD_TOKEN_MEMINDEX
+ return
+
+ def alloc_c4(self, need_size) -> torch.Tensor:
+ raise AssertionError("DeepSeek-V4 c4 uses page-safe allocation; call alloc_c4_pages instead")
+
+ def alloc_c128(self, need_size) -> torch.Tensor:
+ return self.c128_allocator.alloc(need_size)
+
+ def free_c4(self, free_index) -> None:
+ raise AssertionError("DeepSeek-V4 c4 uses page live-count release; call evict_c4 instead")
+
+ def free_c128(self, free_index) -> None:
+ self.c128_allocator.free(free_index)
+
+ # ------------------------------------------------------------------ packed codecs (torch reference)
+ # 与 sglang/vllm 的 fp8_ds_mla 字节布局逐位对齐(ue8m0 幂次 scale)。这些 torch 实现是该 ABI 的
+ # 可执行规格(单测 oracle，triton writer 与其逐字节对拍)，不可删除。
+ def _pack_mla_kv(self, kv: torch.Tensor) -> torch.Tensor:
+ kv = kv.reshape(-1, self.mla_head_dim)
+ out = torch.empty((kv.shape[0], self.mla_bytes_per_token), dtype=torch.uint8, device=kv.device)
+ nope = kv[:, : self.mla_nope_dim].float().reshape(-1, self.mla_scale_bytes - 1, self.mla_quant_group_size)
+ scale = torch.clamp(nope.abs().amax(dim=-1) / DSV4_FP8_E4M3_MAX, min=DSV4_FP8_SCALE_MIN)
+ scale_exp = torch.ceil(torch.log2(scale)).to(torch.int32)
+ scale = torch.exp2(scale_exp.float())
+ nope_fp8 = torch.clamp(nope / scale.unsqueeze(-1), -DSV4_FP8_E4M3_MAX, DSV4_FP8_E4M3_MAX).to(
+ torch.float8_e4m3fn
+ )
+ out[:, : self.mla_nope_dim].copy_(nope_fp8.reshape(-1, self.mla_nope_dim).view(dtype=torch.uint8))
+ rope_start = self.mla_nope_dim
+ rope_end = rope_start + self.mla_rope_dim * 2
+ rope = kv[:, self.mla_nope_dim : self.mla_head_dim].contiguous().to(torch.bfloat16)
+ out[:, rope_start:rope_end].copy_(rope.view(dtype=torch.uint8).reshape(-1, self.mla_rope_dim * 2))
+ scale_start = rope_end
+ scale_end = scale_start + self.mla_scale_bytes - 1
+ out[:, scale_start:scale_end].copy_((scale_exp + 127).to(torch.uint8))
+ out[:, scale_end].zero_()
+ return out
+
+ def _unpack_mla_kv(self, packed: torch.Tensor) -> torch.Tensor:
+ packed = packed.reshape(-1, self.mla_bytes_per_token)
+ if packed.shape[0] == 0:
+ return torch.empty((0, self.mla_head_dim), dtype=self.dtype, device=packed.device)
+ nope_fp8 = packed[:, : self.mla_nope_dim].view(dtype=torch.float8_e4m3fn).float()
+ nope_fp8 = nope_fp8.reshape(-1, self.mla_scale_bytes - 1, self.mla_quant_group_size)
+ rope_start = self.mla_nope_dim
+ rope_end = rope_start + self.mla_rope_dim * 2
+ scale_start = rope_end
+ scale_end = scale_start + self.mla_scale_bytes - 1
+ scale_exp = packed[:, scale_start:scale_end].to(torch.int32) - 127
+ scale = torch.exp2(scale_exp.float())
+ nope = (nope_fp8 * scale.reshape(-1, self.mla_scale_bytes - 1, 1)).reshape(-1, self.mla_nope_dim)
+ rope = packed[:, rope_start:rope_end].view(dtype=torch.bfloat16)
+ return torch.cat([nope.to(self.dtype), rope.to(self.dtype)], dim=-1)
+
+ def _pack_indexer_k(self, indexer_k: torch.Tensor) -> torch.Tensor:
+ indexer_k = indexer_k.reshape(-1, self.indexer_head_dim)
+ out = torch.empty(
+ (indexer_k.shape[0], self.indexer_bytes_per_token),
+ dtype=torch.uint8,
+ device=indexer_k.device,
+ )
+ k_float = indexer_k.float()
+ scale = torch.clamp(
+ k_float.abs().amax(dim=-1, keepdim=True) / DSV4_FP8_E4M3_MAX,
+ min=DSV4_FP8_SCALE_MIN,
+ )
+ k_fp8 = torch.clamp(k_float / scale, -DSV4_FP8_E4M3_MAX, DSV4_FP8_E4M3_MAX).to(torch.float8_e4m3fn)
+ out[:, : self.indexer_head_dim].copy_(k_fp8.view(dtype=torch.uint8))
+ out[:, self.indexer_head_dim :].copy_(scale.view(dtype=torch.uint8).reshape(-1, DSV4_INDEXER_SCALE_BYTES))
+ return out
+
+ def _unpack_indexer_k(self, packed: torch.Tensor) -> torch.Tensor:
+ packed = packed.reshape(-1, self.indexer_bytes_per_token)
+ if packed.shape[0] == 0:
+ return torch.empty((0, self.indexer_head_dim), dtype=self.dtype, device=packed.device)
+ k_fp8 = packed[:, : self.indexer_head_dim].view(dtype=torch.float8_e4m3fn).float()
+ scale = packed[:, self.indexer_head_dim :].view(dtype=torch.float32)
+ return (k_fp8 * scale).to(self.dtype)
+
+ # ------------------------------------------------------------------ cache write paths
+ def pack_mla_kv_to_cache(self, layer_index: int, mem_index: torch.Tensor, kv: torch.Tensor):
+ """标准 operator 写入路径。要求本步已对 mem_index 调过 ``alloc_swa``(prep 阶段)；
+ HOLD/padding 槽位映射到 swa HOLD 槽，写入无害。"""
+ if kv.shape[0] == 0:
+ return
+ from lightllm.models.deepseek_v4.triton_kernel.destindex_copy_kv_flashmla_dsv4 import (
+ destindex_copy_kv_flashmla_dsv4,
+ )
+
+ swa_slots = self.full_to_swa_indexs[mem_index.cuda().long().reshape(-1)]
+ destindex_copy_kv_flashmla_dsv4(
+ kv.reshape(-1, self.mla_head_dim),
+ swa_slots,
+ self.swa_pool.get_layer_buffer(layer_index),
+ self.swa_pool.page_size,
+ )
+ return
+
+ def pack_mla_kv_to_cache_fused_norm_rope(
+ self,
+ layer_index: int,
+ mem_index: torch.Tensor,
+ kv: torch.Tensor,
+ kv_weight: torch.Tensor,
+ eps: float,
+ freqs_cis: torch.Tensor,
+ positions: torch.Tensor,
+ ):
+ """同 pack_mla_kv_to_cache，但 rmsnorm + 尾部交错 rope 融合进写入 kernel
+ 并省掉 bf16 kv 中间量。kv 为 wkv 投影原始输出 [T, head_dim+rope_dim]。"""
+ if kv.shape[0] == 0:
+ return
+ from lightllm.models.deepseek_v4.triton_kernel.norm_rope_cuda import (
+ fused_k_norm_rope_flashmla,
+ )
+
+ swa_slots = self.full_to_swa_indexs[mem_index.cuda().long().reshape(-1)]
+ swa_slots = torch.where(swa_slots < 0, torch.full_like(swa_slots, self.swa_pool.HOLD_TOKEN_MEMINDEX), swa_slots)
+ fused_k_norm_rope_flashmla(
+ kv=kv,
+ kv_weight=kv_weight,
+ eps=eps,
+ freqs_cis=freqs_cis,
+ positions=positions,
+ out_loc=swa_slots,
+ kvcache=self.swa_pool.get_layer_buffer(layer_index),
+ page_size=self.swa_pool.page_size,
+ )
+ return
+
+ def pack_compressed_kv_to_cache(self, layer_index: int, slots: torch.Tensor, comp: torch.Tensor):
+ if comp.shape[0] == 0:
+ return
+ from lightllm.models.deepseek_v4.triton_kernel.destindex_copy_kv_flashmla_dsv4 import (
+ destindex_copy_kv_flashmla_dsv4,
+ )
+
+ pool, local_layer = self._pool_and_local_layer(layer_index)
+ destindex_copy_kv_flashmla_dsv4(
+ comp.reshape(-1, self.mla_head_dim),
+ slots.to(comp.device),
+ pool.get_layer_buffer(local_layer),
+ pool.page_size,
+ )
+
+ def pack_indexer_k_to_cache(self, layer_index: int, slots: torch.Tensor, indexer_k: torch.Tensor):
+ if indexer_k.shape[0] == 0:
+ return
+ assert self.compress_rates[layer_index] == 4, "只有 c4(CSA) 层有 indexer-K"
+ from lightllm.models.deepseek_v4.triton_kernel.destindex_copy_indexer_k_dsv4 import (
+ destindex_copy_indexer_k_dsv4,
+ )
+
+ destindex_copy_indexer_k_dsv4(
+ indexer_k.reshape(-1, self.indexer_head_dim),
+ slots.to(indexer_k.device),
+ self.c4_indexer_pool.get_layer_buffer(self.layer_to_c4_idx[layer_index]),
+ self.c4_indexer_pool.page_size,
+ )
+
+ def gather_indexer_k(self, layer_index: int, slots: torch.Tensor) -> torch.Tensor:
+ """反量化 gather c4 indexer-K: slots [N](c4 槽位,HOLD 合法) -> [N, indexer_head_dim] bf16。
+ indexer top-k 打分用(纯张量操作,cuda-graph 安全)。"""
+ assert self.compress_rates[layer_index] == 4, "只有 c4(CSA) 层有 indexer-K"
+ pool = self.c4_indexer_pool
+ flat = pool.get_layer_buffer(self.layer_to_c4_idx[layer_index]).view(-1)
+ data_offsets, scale_offsets = pool._loc_offsets(slots.reshape(-1))
+ data_range = torch.arange(pool.data_bytes_per_token, device=flat.device)
+ scale_range = torch.arange(pool.scale_bytes_per_token, device=flat.device)
+ k_fp8 = flat[data_offsets.unsqueeze(1) + data_range.unsqueeze(0)].view(torch.float8_e4m3fn)
+ scale = flat[scale_offsets.unsqueeze(1) + scale_range.unsqueeze(0)].contiguous().view(torch.float32)
+ return (k_fp8.float() * scale).to(torch.bfloat16)
+
+ # ------------------------------------------------------------------ fenced inherited APIs
+ # kv_buffer 是 page 索引的 uint8 slab，基类按 token 索引读写的接口会静默写坏数据，显式拦截。
+ def get_index_kv_buffer(self, index):
+ raise NotImplementedError("DeepSeek-V4 packed page-slab cache does not support token-indexed kv_buffer io")
+
+ def load_index_kv_buffer(self, index, load_tensor_dict):
+ raise NotImplementedError("DeepSeek-V4 packed page-slab cache does not support token-indexed kv_buffer io")
+
+ def alloc_kv_move_buffer(self, max_req_total_len):
+ raise NotImplementedError("DeepSeek-V4 packed/composite KV transfer is not implemented")
+
+ def alloc_paged_kv_move_buffer(self, page_num, page_size) -> torch.Tensor:
+ raise NotImplementedError("DeepSeek-V4 packed/composite KV transfer is not implemented")
+
+ def write_mem_to_page_kv_move_buffer(self, *args, **kwargs):
+ raise NotImplementedError("DeepSeek-V4 packed/composite KV transfer is not implemented")
+
+ def read_page_kv_move_buffer_to_mem(self, *args, **kwargs):
+ raise NotImplementedError("DeepSeek-V4 packed/composite KV transfer is not implemented")
+
+ def send_to_decode_node(self, *args, **kwargs):
+ raise NotImplementedError("DeepSeek-V4 packed/composite KV transfer is not implemented")
+
+ def receive_from_prefill_node(self, *args, **kwargs):
+ raise NotImplementedError("DeepSeek-V4 packed/composite KV transfer is not implemented")
+
+ def send_to_decode_node_p2p(self, *args, **kwargs):
+ raise NotImplementedError("DeepSeek-V4 packed/composite KV transfer is not implemented")
+
+ def receive_from_prefill_node_p2p(self, *args, **kwargs):
+ raise NotImplementedError("DeepSeek-V4 packed/composite KV transfer is not implemented")
diff --git a/lightllm/common/kv_cache_mem_manager/operator/__init__.py b/lightllm/common/kv_cache_mem_manager/operator/__init__.py
index 85c37ad39b..442c2e300e 100644
--- a/lightllm/common/kv_cache_mem_manager/operator/__init__.py
+++ b/lightllm/common/kv_cache_mem_manager/operator/__init__.py
@@ -5,6 +5,7 @@
from .deepseek import (
Deepseek2MemOperator,
Deepseek3_2MemOperator,
+ DeepseekV4MemOperator,
FP8PerTokenGroupQuantDeepseek3_2MemOperator,
)
from .fp8_quant import (
diff --git a/lightllm/common/kv_cache_mem_manager/operator/deepseek.py b/lightllm/common/kv_cache_mem_manager/operator/deepseek.py
index 6e05b96e10..0725ce9b93 100644
--- a/lightllm/common/kv_cache_mem_manager/operator/deepseek.py
+++ b/lightllm/common/kv_cache_mem_manager/operator/deepseek.py
@@ -8,7 +8,9 @@
class Deepseek2MemOperator(NormalMemOperator):
def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv: torch.Tensor):
- from lightllm.common.kv_cache_mem_manager.deepseek2_mem_manager import Deepseek2MemoryManager
+ from lightllm.common.kv_cache_mem_manager.deepseek2_mem_manager import (
+ Deepseek2MemoryManager,
+ )
mem_manager: Deepseek2MemoryManager = self.mem_manager
@@ -30,7 +32,9 @@ def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv:
class Deepseek3_2MemOperator(Deepseek2MemOperator):
def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv: torch.Tensor):
- from lightllm.common.kv_cache_mem_manager.deepseek3_2mem_manager import Deepseek3_2MemoryManager
+ from lightllm.common.kv_cache_mem_manager.deepseek3_2mem_manager import (
+ Deepseek3_2MemoryManager,
+ )
mem_manager: Deepseek3_2MemoryManager = self.mem_manager
from ...basemodel.triton_kernel.kv_copy.mla_copy_kv import destindex_copy_kv
@@ -78,3 +82,14 @@ def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv:
o_rope,
)
return
+
+
+class DeepseekV4MemOperator(BaseMemManagerOperator):
+ def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv: torch.Tensor):
+ from lightllm.common.kv_cache_mem_manager.deepseek4_mem_manager import (
+ DeepseekV4MemoryManager,
+ )
+
+ mem_manager: DeepseekV4MemoryManager = self.mem_manager
+ mem_manager.pack_mla_kv_to_cache(layer_index, mem_index, kv)
+ return
diff --git a/lightllm/common/quantization/__init__.py b/lightllm/common/quantization/__init__.py
index cd534d53ec..4db55e1555 100644
--- a/lightllm/common/quantization/__init__.py
+++ b/lightllm/common/quantization/__init__.py
@@ -14,6 +14,7 @@
EXPERT_DTYPE_TO_QUANT_TYPE = {
"fp8": "deepgemm-fp8w8a8-b128",
"fp4": "deepgemm-fp4fp8-b32",
+ "mxfp4": "marlin-mxfp4w4a16-b32",
}
SUPPORTED_EXPERT_DTYPES = tuple(EXPERT_DTYPE_TO_QUANT_TYPE)
@@ -64,10 +65,13 @@ def _mapping_quant_method(self):
logger.info(f"select fp8w8a8-b128 quant way: {self.quant_type}")
# fp8 量化下，部分 MoE 模型（如 DeepSeek-V4），可以单独声明 expert 权重精度，
- # 按其值给 fused_moe 选用对应的 deepgemm 量化方法。
+ # 按其值给 fused_moe 选用对应的量化方法。
expert_dtype = self.expert_dtype or self.network_config_.get("expert_dtype", None)
if expert_dtype is None:
return
+ # DeepSeek-V4 的 fp4 发布版自带预打包 MXFP4 专家。
+ if expert_dtype == "fp4" and self.network_config_.get("model_type") == "deepseek_v4":
+ expert_dtype = "mxfp4"
target = self._get_expert_quant_type(expert_dtype)
for layer_num in range(self.layer_num):
if self.expert_dtype is not None:
diff --git a/lightllm/common/quantization/deepgemm.py b/lightllm/common/quantization/deepgemm.py
index ec1ee90fd4..bedf22ee95 100644
--- a/lightllm/common/quantization/deepgemm.py
+++ b/lightllm/common/quantization/deepgemm.py
@@ -198,6 +198,144 @@ def _create_weight(
return mm_param, mm_param_list
+@QUANTMETHODS.register(["marlin-mxfp4w4a16-b32"], platform="cuda")
+class MXFP4MoEQuantizationMethod(QuantizationMethod):
+ def __init__(self):
+ super().__init__()
+ self.block_size = 32
+ self.weight_suffix = "weight"
+ self.weight_zero_point_suffix = None
+ self.weight_scale_suffix = "scale"
+ self.has_weight_scale = True
+ self.has_weight_zero_point = False
+
+ @property
+ def method_name(self):
+ return "marlin-mxfp4w4a16-b32"
+
+ def quantize(self, weight: torch.Tensor, output: WeightPack):
+ raise NotImplementedError("marlin-mxfp4w4a16-b32 only loads pre-packed MXFP4 expert weights")
+
+ def apply(
+ self,
+ input_tensor: torch.Tensor,
+ weight_pack: "WeightPack",
+ out: Optional[torch.Tensor] = None,
+ workspace: Optional[torch.Tensor] = None,
+ use_custom_tensor_mananger: bool = True,
+ bias: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ raise NotImplementedError("marlin-mxfp4w4a16-b32 is only implemented for fused MoE expert weights")
+
+ def _probe_marlin_layout(self, size_n: int, size_k: int, dtype: torch.dtype, device_id: int):
+ """用零输入走一遍真实的 per-expert repack 路径,探出 marlin 终态布局的形状与类型。
+ 只调用 finalize 同款的 vllm 函数,不复刻其内部公式,杜绝形状漂移。结果按维度缓存
+ (各 MoE 层同维,全程只探两次: w13 一次、w2 一次)。"""
+ cache_key = (size_n, size_k, dtype)
+ cache = getattr(self, "_marlin_layout_cache", None)
+ if cache is None:
+ cache = self._marlin_layout_cache = {}
+ if cache_key in cache:
+ return cache[cache_key]
+
+ import vllm._custom_ops as ops
+ from vllm.model_executor.layers.quantization.utils.marlin_utils import (
+ get_marlin_input_dtype,
+ marlin_permute_scales,
+ )
+ from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
+ mxfp4_marlin_process_scales,
+ )
+
+ input_dtype = get_marlin_input_dtype()
+ is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1
+ device = f"cuda:{device_id}"
+ qweight = torch.zeros((size_n, size_k // 2), dtype=torch.int8, device=device).view(torch.int32).T.contiguous()
+ marlin_qweight = ops.gptq_marlin_repack(
+ b_q_weight=qweight,
+ perm=torch.empty(0, dtype=torch.int, device=device),
+ size_k=size_k,
+ size_n=size_n,
+ num_bits=4,
+ is_a_8bit=is_a_8bit,
+ )
+ scale = torch.zeros((size_k // self.block_size, size_n), dtype=dtype, device=device)
+ marlin_scale = marlin_permute_scales(
+ s=scale, size_k=size_k, size_n=size_n, group_size=self.block_size, is_a_8bit=is_a_8bit
+ )
+ marlin_scale = mxfp4_marlin_process_scales(marlin_scale, input_dtype=input_dtype)
+ layout = (
+ (tuple(marlin_qweight.shape), marlin_qweight.dtype),
+ (tuple(marlin_scale.shape), marlin_scale.dtype),
+ )
+ cache[cache_key] = layout
+ return layout
+
+ def _create_weight(
+ self, out_dims: Union[int, List[int]], in_dim: int, dtype: torch.dtype, device_id: int, num_experts: int = 1
+ ) -> Tuple[WeightPack, List[WeightPack]]:
+ out_dim = sum(out_dims) if isinstance(out_dims, list) else out_dims
+ assert in_dim % self.block_size == 0, "MXFP4 scale dimension must be divisible by block_size"
+ expert_prefix = (num_experts,) if num_experts > 1 else ()
+ # CPU 暂存区: load_hf_weights 灌入原始预打包 MXFP4,finalize 时 repack 进 CUDA 终态。
+ weight = torch.empty(expert_prefix + (out_dim, in_dim // 2), dtype=torch.int8, device="cpu")
+ weight_scale = torch.empty(
+ expert_prefix + (out_dim, in_dim // self.block_size), dtype=torch.float8_e8m0fnu, device="cpu"
+ )
+ mm_param = WeightPack(weight=weight, weight_scale=weight_scale)
+ # CUDA 终态(marlin 布局)在构造期物化,使 mem manager 的 profile 看到真实权重占用
+ # ("构造即分配、load 只灌数"的框架契约,与其它 quant 方法一致;惰性到 finalize 才
+ # 进卡会让空卡 profile 把 kv 池撑到挤爆权重加载)。finalize 时 repack 结果拷入。
+ (w_shape, w_dtype), (s_shape, s_dtype) = self._probe_marlin_layout(out_dim, in_dim, dtype, device_id)
+ mm_param.marlin_weight = torch.empty((num_experts,) + w_shape, dtype=w_dtype, device=f"cuda:{device_id}")
+ mm_param.marlin_weight_scale = torch.empty((num_experts,) + s_shape, dtype=s_dtype, device=f"cuda:{device_id}")
+ mm_param_list = self._split_weight_pack(
+ mm_param,
+ weight_out_dims=out_dims,
+ weight_split_dim=-2,
+ weight_scale_out_dims=out_dims,
+ weight_scale_split_dim=-2,
+ )
+ return mm_param, mm_param_list
+
+ def finalize_moe_weight(self, moe_weight):
+ try:
+ from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
+ prepare_moe_mxfp4_layer_for_marlin,
+ )
+ except Exception as e:
+ raise RuntimeError(f"marlin-mxfp4w4a16-b32 requires vLLM MXFP4 packing utilities, error={repr(e)}") from e
+
+ class _MXFP4Layer:
+ pass
+
+ device = torch.device("cuda", moe_weight.device_id_)
+ layer = _MXFP4Layer()
+ layer.params_dtype = moe_weight.data_type_
+ w13 = moe_weight.w13.weight.view(torch.uint8).to(device=device, non_blocking=True).contiguous()
+ w2 = moe_weight.w2.weight.view(torch.uint8).to(device=device, non_blocking=True).contiguous()
+ w13_scale = moe_weight.w13.weight_scale.to(device=device, non_blocking=True).contiguous()
+ w2_scale = moe_weight.w2.weight_scale.to(device=device, non_blocking=True).contiguous()
+ (
+ w13_new,
+ w2_new,
+ w13_scale_new,
+ w2_scale_new,
+ _,
+ _,
+ ) = prepare_moe_mxfp4_layer_for_marlin(layer, w13, w2, w13_scale, w2_scale, None, None)
+ # repack 结果拷入构造期预分配的 marlin 终态 buffer(与 AWQ marlin 路径同形态),
+ # CPU 暂存与 repack 临时随引用释放;shape 失配会在 copy_ 处显式报错(探针保证一致)。
+ moe_weight.w13.marlin_weight.copy_(w13_new)
+ moe_weight.w13.marlin_weight_scale.copy_(w13_scale_new)
+ moe_weight.w2.marlin_weight.copy_(w2_new)
+ moe_weight.w2.marlin_weight_scale.copy_(w2_scale_new)
+ moe_weight.w13.weight = moe_weight.w13.marlin_weight
+ moe_weight.w13.weight_scale = moe_weight.w13.marlin_weight_scale
+ moe_weight.w2.weight = moe_weight.w2.marlin_weight
+ moe_weight.w2.weight_scale = moe_weight.w2.marlin_weight_scale
+
+
def _deepgemm_fp8_nt(a_tuple, b_tuple, out):
if HAS_DEEPGEMM:
if hasattr(deep_gemm, "gemm_fp8_fp8_bf16_nt"):
diff --git a/lightllm/common/quantization/no_quant.py b/lightllm/common/quantization/no_quant.py
index fa926ad6f0..21e7101d4d 100644
--- a/lightllm/common/quantization/no_quant.py
+++ b/lightllm/common/quantization/no_quant.py
@@ -18,19 +18,26 @@ def apply(
workspace: Optional[torch.Tensor] = None,
use_custom_tensor_mananger: bool = True,
bias: Optional[torch.Tensor] = None,
+ out_dtype: Optional[torch.dtype] = None,
) -> torch.Tensor:
from lightllm.common.basemodel.layer_infer.cache_tensor_manager import g_cache_manager
weight = weight_pack.weight.t()
+ if out_dtype is not None and bias is not None:
+ raise NotImplementedError("out_dtype is only supported for bias-free no-quant mm")
if out is None:
shape = (input_tensor.shape[0], weight.shape[1])
- dtype = input_tensor.dtype
+ dtype = out_dtype if out_dtype is not None else input_tensor.dtype
device = input_tensor.device
if use_custom_tensor_mananger:
out = g_cache_manager.alloc_tensor(shape, dtype, device=device)
else:
out = torch.empty(shape, dtype=dtype, device=device)
+ elif out_dtype is not None and out.dtype != out_dtype:
+ raise ValueError(f"out dtype {out.dtype} does not match requested out_dtype {out_dtype}")
if bias is None:
+ if out_dtype is not None:
+ return torch.mm(input_tensor, weight, out=out, out_dtype=out_dtype)
return torch.mm(input_tensor, weight, out=out)
return torch.addmm(bias, input_tensor, weight, out=out)
diff --git a/lightllm/common/quantization/quantize_method.py b/lightllm/common/quantization/quantize_method.py
index 95d8d806f9..d3f251ec84 100644
--- a/lightllm/common/quantization/quantize_method.py
+++ b/lightllm/common/quantization/quantize_method.py
@@ -55,6 +55,7 @@ def apply(
workspace: Optional[torch.Tensor] = None,
use_custom_tensor_mananger: bool = True,
bias: Optional[torch.Tensor] = None,
+ out_dtype: Optional[torch.dtype] = None,
) -> torch.Tensor:
pass
diff --git a/lightllm/common/req_manager.py b/lightllm/common/req_manager.py
index 01e9c4ad35..fe49e9a095 100644
--- a/lightllm/common/req_manager.py
+++ b/lightllm/common/req_manager.py
@@ -1,17 +1,26 @@
import torch
import collections
+from dataclasses import dataclass
from lightllm.common.linear_att_cache_manager.config_objs import LinearAttCacheConfig
from lightllm.utils.log_utils import init_logger
-from .kv_cache_mem_manager import MemoryManager
+from .kv_cache_mem_manager import MemoryManager, DeepseekV4MemoryManager
from typing import List, Optional, TYPE_CHECKING
from lightllm.common.basemodel.triton_kernel.gen_sampling_params import token_id_counter
-from lightllm.common.basemodel.triton_kernel.gen_sampling_params import update_req_to_token_id_counter
+from lightllm.common.basemodel.triton_kernel.gen_sampling_params import (
+ update_req_to_token_id_counter,
+)
from lightllm.utils.envs_utils import enable_env_vars, get_env_start_args
from lightllm.utils.config_utils import get_vocab_size
from lightllm.server.router.model_infer.pin_mem_manager import g_pin_mem_manager
from lightllm.common.linear_att_cache_manager.layer_cache import LayerCache
-from lightllm.common.linear_att_cache_manager.linear_att_buffer_manager import LinearAttCacheManager
+from lightllm.common.linear_att_cache_manager.linear_att_buffer_manager import (
+ LinearAttCacheManager,
+)
+from lightllm.common.kv_cache_mem_manager.deepseek4_mem_manager import (
+ DSV4_C4_PAGE_SIZE,
+ DSV4_PROMPT_CACHE_PAGE_SIZE,
+)
if TYPE_CHECKING:
from lightllm.server.router.model_infer.infer_batch import InferReq
@@ -19,6 +28,61 @@
logger = init_logger(__name__)
+@dataclass
+class DeepseekV4PromptCachePayload:
+ """prompt cache 载荷: swa 按页有效性 bitmap 和最后有效页。
+
+ 槽位与 compressor 状态都不进载荷: full_to_swa/full_to_c4/full_to_c128 以 full token 槽位
+ 为键(radix 持有 full 槽 ⇒ 映射行存活,free 级联回收);c4/c128 compressor 状态以 swa
+ 页派生寻址(随 swa 页生灭,命中零拷贝续算)。prompt cache 对齐到 256 token,
+ 避免共享前缀停在 c4 物理页中间。
+
+ * ``swa_page_valid``: cpu bool [cache_len // page]，插入时按当下 full_to_swa 映射写定
+ (页内 token 映射全有效才为 True)。匹配层据此把命中裁剪到"结尾页有效"的 page 边界,
+ swa 压力阀回收节点页时清零。"""
+
+ cache_len: int
+ swa_page_valid: Optional[torch.Tensor] = None
+ swa_last_valid_page: int = -1
+
+ def refresh_swa_last_valid_page(self) -> None:
+ if self.swa_page_valid is None:
+ self.swa_last_valid_page = -1
+ return
+ valid_idx = torch.nonzero(self.swa_page_valid).flatten()
+ self.swa_last_valid_page = -1 if valid_idx.numel() == 0 else int(valid_idx[-1].item())
+ return
+
+ def valid_match_length(self, natural_len: int, page: int) -> int:
+ if self.swa_last_valid_page < 0:
+ return 0
+ return (int(self.swa_last_valid_page) + 1) * page
+
+
+class DeepseekV4PromptCacheValueOps:
+ def __init__(self, req_manager: "DeepseekV4ReqManager"):
+ self.req_manager = req_manager
+
+ def slice(self, payload: DeepseekV4PromptCachePayload, start: int, end: int):
+ return self.req_manager.slice_prompt_cache_payload(payload, start, end)
+
+ def concat(self, payloads: List[DeepseekV4PromptCachePayload]):
+ return self.req_manager.concat_prompt_cache_payloads(payloads)
+
+ def free(self, payload: DeepseekV4PromptCachePayload):
+ # 槽位资源全部由 mem_manager.free(full_slots) 级联回收，载荷本身没有需要释放的资源。
+ return
+
+ def valid_match_length(self, payload: Optional[DeepseekV4PromptCachePayload], natural_len: int) -> int:
+ """radix 匹配裁剪: 返回 <= natural_len 的最大 prompt-cache 边界 L'，使结尾页有效。
+
+ 有效性可能非单调(owner 生前从左驱逐、后续阀从尾回收)，中段 invalid 页不挡更
+ 靠后的有效命中(注意力只回看最后一个窗口)。"""
+ if payload is None:
+ return 0
+ return payload.valid_match_length(natural_len, self.req_manager.get_prompt_cache_page_size())
+
+
class _ReqNode:
def __init__(self, index):
self.index = index
@@ -299,3 +363,568 @@ def copy_small_page_buffer_to_linear_att_state(
self.req_to_conv_state.buffer[:, dest_req_idx, ...] = conv_state
self.req_to_ssm_state.buffer[:, dest_req_idx, ...] = ssm_state
return
+
+
+class DeepseekV4ReqManager(ReqManager):
+ """DeepSeek-V4 的请求级管理。
+
+ 在基类 ReqManager 之上补 V4 专有的 per-request 结构。该对象在 mem manager profile 前创建，
+ 所以初始化只依赖 config 派生出的 compress_rates/head_dim/indexer_head_dim/sliding_window；
+ 真实 mem_manager 会在 `_init_mem_manager()` 后通过 `bind_mem_manager()` 接入。
+
+ * 压缩槽位不在本类: ``full_to_c4/c128_indexs``(mem manager)以组末 token 的 full 槽位为键。
+ 本类只负责 prep 阶段的分配与 scatter(``prepare_prefill`` /
+ ``prepare_decode_compress_slots``)——必须先于 attention metadata 构建/图捕获;
+ 条目内容由 layer-infer 的 compressor 前向写入。
+ * compressor 在途状态不在本类: c4/c128 都在 mem manager 的 swa 页派生池,
+ 随页生灭,命中零拷贝续算。
+ * SWA 槽位分配/出窗回收(``prepare_prefill_swa`` / ``prepare_decode_swa``): 每步 prep 阶段
+ 为新 token 调 mem_manager.alloc_swa，并按 per-req 水位线(``_swa_evict_marks``)惰性回收
+ 已出窗位置的 swa 槽。水位线首次置为该请求首个 chunk 的 ready_cache_len(radix 共享前缀
+ 的边界)，因此共享前缀的 swa 槽永远不会被本请求回收(归 radix 经 mem_manager.free 级联释放)。
+ """
+
+ def __init__(
+ self,
+ max_request_num,
+ max_sequence_length,
+ mem_manager: Optional[DeepseekV4MemoryManager] = None,
+ compress_rates: Optional[List[int]] = None,
+ head_dim: Optional[int] = None,
+ indexer_head_dim: Optional[int] = None,
+ sliding_window: Optional[int] = None,
+ ):
+ super().__init__(max_request_num, max_sequence_length, mem_manager)
+
+ self.sliding_window = sliding_window
+ # 出窗回收水位线: -1 表示该 req 尚未见过 prefill chunk(首个 chunk 的 ready_cache_len
+ # 即共享前缀边界，作为永不下探的回收下界)。
+ self._swa_evict_marks = [-1 for _ in range(max_request_num + 1)]
+ self.compress_rates = list(compress_rates)
+ self.n_c4 = sum(1 for r in self.compress_rates if r == 4)
+ self.n_c128 = sum(1 for r in self.compress_rates if r == 128)
+ self.head_dim = head_dim
+ self.indexer_head_dim = indexer_head_dim
+ self.layer_to_c4_idx = {}
+ self.layer_to_c128_idx = {}
+ self.mem_manager = mem_manager
+ c4 = c128 = 0
+ for lid, r in enumerate(self.compress_rates):
+ if r == 4:
+ self.layer_to_c4_idx[lid] = c4
+ c4 += 1
+ elif r == 128:
+ self.layer_to_c128_idx[lid] = c128
+ c128 += 1
+
+ return
+
+ # ------------------------------------------------------------------ swa slot prep (per step)
+ def _swa_retain_len(self) -> int:
+ """出窗回收的保留长度 = window + 一个 radix 页。
+
+ 多留一页使「最近一个完成的 prompt-cache 边界」的结尾页恒驻留: 若回收只留 window，
+ 则任何非对齐时刻该边界的结尾页都已被部分回收，插入门会把所有插入裁到 0。
+ V4 prompt-cache 页取 256 token，正好覆盖一个 c4 物理页对应的 token 范围。"""
+ return int(self.sliding_window) + self.get_prompt_cache_page_size()
+
+ def _align_swa_evict_frontier(self, raw_frontier: int) -> int:
+ """SWA 回收水位线按 prompt-cache 页向下对齐。
+
+ bitmap 的有效性是 prompt-cache page 粒度；若水位线切进页面中间，该页会被判为
+ invalid，即使靠近命中边界的窗口实际仍完整驻留。"""
+ page = self.get_prompt_cache_page_size()
+ raw_frontier = max(0, int(raw_frontier))
+ return raw_frontier // page * page
+
+ def prepare_prefill_swa(
+ self,
+ req_list: List[int],
+ ready_list: List[int],
+ seq_list: List[int],
+ mem_indexes: torch.Tensor,
+ ) -> None:
+ """prefill prep: 为本 chunk 全部新 token(位置 [ready, seq))分配位置对齐的 swa 槽，
+ 并回收已出窗位置的槽。
+
+ 本 chunk 起点 L = ready_cache_len，首个新 token(位置 L)的窗口是 [L-W+1, L]；回收
+ 边界再额外保留一个 radix 页(_swa_retain_len)，即位置 < L-retain+1。先回收再分配。
+ 当前 chunk 的 full slots 直接使用 generic preprocess 分配的 mem_indexes，因而可以
+ 在通用 req_to_token scatter 之前执行。"""
+ self.mem_manager: DeepseekV4MemoryManager
+ if self.sliding_window is not None:
+ retain = self._swa_retain_len()
+ evict_slots = []
+ for req_idx, ready_len in zip(req_list, ready_list):
+ if req_idx == self.HOLD_REQUEST_ID:
+ continue
+ mark = self._swa_evict_marks[req_idx]
+ if mark < 0:
+ # 首个 chunk: [0, ready_len) 是 radix 共享前缀，其 swa 槽归 radix 所有，不可回收。
+ self._swa_evict_marks[req_idx] = self._align_swa_evict_frontier(ready_len)
+ continue
+ evict_end = self._align_swa_evict_frontier(ready_len - retain + 1)
+ if evict_end > mark:
+ evict_slots.append(self.req_to_token_indexs[req_idx, mark:evict_end])
+ self._swa_evict_marks[req_idx] = evict_end
+ if evict_slots:
+ self.mem_manager.evict_swa(torch.cat(evict_slots))
+ self.mem_manager.alloc_swa_prefill(
+ mem_indexes,
+ self.req_to_token_indexs,
+ req_list=req_list,
+ ready_list=ready_list,
+ seq_list=seq_list,
+ )
+ return
+
+ def prepare_decode(
+ self,
+ b_req_idx_cpu,
+ b_seq_len_cpu,
+ b_mtp_index_cpu,
+ mem_indexes,
+ mtp_decode_slot_prepare_indices,
+ prepare_compress_slots=True,
+ ):
+ """decode 每步槽位 prep。在 BaseModel 的通用 req scatter 与 attention metadata
+ 构建前调用；DeepSeek-V4 MTP draft layer 只需要 SWA 槽位。"""
+ max_mtp_index = int(b_mtp_index_cpu.max().item())
+ if mtp_decode_slot_prepare_indices is None:
+ steps = range(max_mtp_index + 1)
+ else:
+ steps = mtp_decode_slot_prepare_indices
+
+ batch_size = b_mtp_index_cpu.shape[0]
+ slots_per_req = max_mtp_index + 1
+ assert batch_size % slots_per_req == 0
+ req_list = b_req_idx_cpu.tolist()
+ seq_list = b_seq_len_cpu.tolist()
+ mem_indexes_by_req = mem_indexes.reshape(-1, slots_per_req)
+ for step in steps:
+ step_req_list = req_list[step::slots_per_req]
+ step_seq_list = seq_list[step::slots_per_req]
+ self.prepare_decode_swa(
+ step_req_list,
+ step_seq_list,
+ mem_indexes_by_req[:, step],
+ prev_mem_indexes=mem_indexes_by_req[:, step - 1] if step > 0 else None,
+ )
+ if prepare_compress_slots:
+ self.prepare_decode_compress_slots(
+ step_req_list,
+ step_seq_list,
+ mem_indexes_by_req[:, step],
+ prev_group_end_mem_indexes=mem_indexes_by_req[:, step - 4] if step >= 4 else None,
+ )
+ return
+
+ def prepare_prefill(
+ self,
+ b_req_idx_cpu: torch.Tensor,
+ b_ready_cache_len_cpu: torch.Tensor,
+ b_seq_len_cpu: torch.Tensor,
+ mem_indexes: torch.Tensor,
+ ) -> None:
+ """prefill 槽位 prep: 直接消费 generic preprocess 分配的 full slots，在
+ BaseModel 的通用 req scatter 与 attention metadata 构建之前完成。"""
+ req_list = b_req_idx_cpu.tolist()
+ ready_list = b_ready_cache_len_cpu.tolist()
+ seq_list = b_seq_len_cpu.tolist()
+ mem_indexes = mem_indexes.reshape(-1)
+ self.prepare_prefill_swa(
+ req_list=req_list,
+ ready_list=ready_list,
+ seq_list=seq_list,
+ mem_indexes=mem_indexes,
+ )
+ self.prepare_prefill_compress_slots(
+ req_list=req_list,
+ ready_list=ready_list,
+ seq_list=seq_list,
+ mem_indexes=mem_indexes,
+ )
+ return
+
+ def prepare_decode_swa(
+ self,
+ req_list: List[int],
+ seq_list: List[int],
+ mem_indexes: torch.Tensor,
+ prev_mem_indexes: Optional[torch.Tensor] = None,
+ ) -> None:
+ """decode prep: 回收出窗槽并为本步新 token 分配位置对齐的 swa 槽。当前 query 位置
+ seq_len-1 的窗口是 [seq_len-W, seq_len-1]；回收边界额外保留一个 radix 页
+ (_swa_retain_len)，即位置 < seq_len-retain。先回收再分配。
+ seq_len/req_idx 从 CPU 镜像读(host 算术,无 D2H);水位线 _swa_evict_marks 仍是 host 状态。"""
+ assert self.mem_manager is not None
+ if self.sliding_window is not None:
+ retain = self._swa_retain_len()
+ evict_slots = []
+ for req_idx, seq_len in zip(req_list, seq_list):
+ if req_idx == self.HOLD_REQUEST_ID:
+ continue
+ mark = self._swa_evict_marks[req_idx]
+ if mark < 0:
+ # 未经过 prefill prep 的保守路径: 不回收旧位置，仅推进水位线。
+ self._swa_evict_marks[req_idx] = self._align_swa_evict_frontier(seq_len - retain)
+ continue
+ evict_end = self._align_swa_evict_frontier(seq_len - retain)
+ if evict_end > mark:
+ evict_slots.append(self.req_to_token_indexs[req_idx, mark:evict_end])
+ self._swa_evict_marks[req_idx] = evict_end
+ if evict_slots:
+ self.mem_manager.evict_swa(torch.cat(evict_slots))
+ if prev_mem_indexes is None:
+ prev_meta = g_pin_mem_manager.gen_from_list(
+ key="dsv4_swa_decode_prev",
+ data=[x for req_idx, seq_len in zip(req_list, seq_list) for x in (req_idx, seq_len - 2)],
+ dtype=torch.int64,
+ ).to(self.req_to_token_indexs.device, non_blocking=True)
+ prev_meta = prev_meta.view(-1, 2)
+ prev_mem_indexes = self.req_to_token_indexs[prev_meta[:, 0], prev_meta[:, 1]]
+ self.mem_manager.alloc_swa_decode(
+ req_list,
+ seq_list,
+ mem_indexes,
+ prev_mem_indexes,
+ )
+ return
+
+ def init_compress_state(self, req_idx: int):
+ """新请求开始时重置 runtime 水位线(对应 mamba 的 init_linear_att_state 调用点)。
+
+ c4/c128 compressor state 都随 swa 页寻址,由内核按组覆写;请求复用时不做 per-req 清零。"""
+ self.clear_runtime_state(req_idx)
+ return
+
+ # ------------------------------------------------------------------ compress slot prep (per step)
+ def _register_c4_slots(self, full_slots: torch.Tensor, slots: torch.Tensor) -> None:
+ """写入 full->c4 槽映射并按页累加存活计数。"""
+ self.mem_manager.full_to_c4_indexs[full_slots] = slots
+ self.mem_manager.count_c4_slots(slots, 1)
+
+ def _scatter_c4_prefill_slots_batched(self, req_list, ready_list, seq_list, mem_indexes) -> None:
+ """Batch c4 prefill scatter from the generic preprocess full-slot layout.
+
+ Each group's end token is in the current chunk, so its full slot is addressed directly in
+ mem_indexes. Only a mid-page continuation reads the previous group's old req-table entry.
+ New full slots guarantee a fresh mapping; no GPU-to-CPU idempotency check is needed."""
+ page = DSV4_C4_PAGE_SIZE
+ mapping = self.mem_manager.full_to_c4_indexs
+ device = mapping.device
+
+ plan = []
+ mem_offset = 0
+ for req_idx, ready_len, seq_len in zip(req_list, ready_list, seq_list):
+ q_len = seq_len - ready_len
+ if req_idx == self.HOLD_REQUEST_ID:
+ mem_offset += q_len
+ continue
+ first, last = ready_len // 4, seq_len // 4
+ if last <= first:
+ mem_offset += q_len
+ continue
+ plan.append((req_idx, ready_len, mem_offset, first, last))
+ mem_offset += q_len
+ if not plan:
+ return
+
+ def to_cuda_long(key, data):
+ return g_pin_mem_manager.gen_from_list(key=key, data=data, dtype=torch.int64).to(device, non_blocking=True)
+
+ reqs, readies, mem_offsets, firsts, lasts = zip(*plan)
+ counts = [last - first for first, last in zip(firsts, lasts)]
+ first_pages = [first // page for first in firsts]
+ page_counts = [((last - 1) // page) - fp + 1 for last, fp in zip(lasts, first_pages)]
+ page_offsets, total_pages = [], 0
+ for n_pages in page_counts:
+ page_offsets.append(total_pages)
+ total_pages += n_pages
+ total_entries = sum(counts)
+ cont = [(off, req, first) for off, req, first in zip(page_offsets, reqs, firsts) if first % page != 0]
+ self._realize_c4_pages(total_pages - len(cont))
+
+ # One pinned H2D copy for all per-request metadata, then per-entry ragged expansion.
+ meta = to_cuda_long(
+ "dsv4_c4_prefill_meta",
+ [x for row in zip(readies, mem_offsets, firsts, first_pages, counts, page_offsets) for x in row],
+ ).view(-1, 6)
+ readies_t, mem_offsets_t, firsts_t, first_pages_t, counts_t, page_offsets_t = meta.unbind(1)
+ seg = torch.repeat_interleave(torch.arange(len(plan), device=device), counts_t, output_size=total_entries)
+ seg_starts = counts_t.cumsum(0) - counts_t
+ entries = firsts_t[seg] + torch.arange(total_entries, device=device) - seg_starts[seg]
+ full_offsets = mem_offsets_t[seg] + entries * 4 + 3 - readies_t[seg]
+ full_slots = mem_indexes.reshape(-1)[full_offsets]
+
+ # physical base per logical page: fresh pages from one alloc; mid-page continuations read prev
+ if not cont:
+ page_bases = self.mem_manager.alloc_c4_pages(total_pages).to(device, non_blocking=True) * page
+ else:
+ page_bases = torch.empty(total_pages, dtype=torch.int32, device=device)
+ new_pos = [
+ pos
+ for off, n_pages, first in zip(page_offsets, page_counts, firsts)
+ for pos in range(off + (first % page != 0), off + n_pages)
+ ]
+ if new_pos:
+ new_pos_t = to_cuda_long("dsv4_c4_prefill_new_pos", new_pos)
+ page_bases[new_pos_t] = (
+ self.mem_manager.alloc_c4_pages(len(new_pos)).to(device, non_blocking=True) * page
+ )
+ cont_t = to_cuda_long("dsv4_c4_prefill_cont", [x for row in cont for x in row]).view(-1, 3)
+ prev_slot = mapping[self.req_to_token_indexs[cont_t[:, 1], cont_t[:, 2] * 4 - 1]]
+ cont_off = ((cont_t[:, 2] - 1) % page).to(torch.int32)
+ page_bases[cont_t[:, 0]] = prev_slot - cont_off
+
+ page_idx = page_offsets_t[seg] + torch.div(entries, page, rounding_mode="floor") - first_pages_t[seg]
+ slots = page_bases[page_idx] + (entries % page).to(torch.int32)
+ self._register_c4_slots(full_slots, slots)
+ return
+
+ def _scatter_c4_decode_slots(
+ self,
+ req_list,
+ seq_list,
+ mem_indexes: torch.Tensor,
+ prev_group_end_mem_indexes: Optional[torch.Tensor] = None,
+ ) -> None:
+ page = DSV4_C4_PAGE_SIZE
+ mapping = self.mem_manager.full_to_c4_indexs
+ mem_indexes = mem_indexes.reshape(-1)
+
+ cont_rows, cont_prev_pos = [], []
+ new_rows = []
+ for i, (req_idx, seq_len) in enumerate(zip(req_list, seq_list)):
+ if req_idx == self.HOLD_REQUEST_ID or seq_len <= 0 or seq_len % 4 != 0:
+ continue
+ entry = seq_len // 4 - 1
+ offset = entry % page
+ if offset == 0:
+ new_rows.append(i)
+ else:
+ cont_rows.append(i)
+ cont_prev_pos.append(entry * 4 - 1)
+
+ if cont_rows:
+ if prev_group_end_mem_indexes is None:
+ prev_meta = g_pin_mem_manager.gen_from_list(
+ key="dsv4_c4_decode_prev",
+ data=[x for row in zip([req_list[i] for i in cont_rows], cont_prev_pos) for x in row],
+ dtype=torch.int64,
+ ).to(mapping.device, non_blocking=True)
+ prev_meta = prev_meta.view(-1, 2)
+ prev_full = self.req_to_token_indexs[prev_meta[:, 0], prev_meta[:, 1]]
+ else:
+ prev_full = prev_group_end_mem_indexes.reshape(-1)[cont_rows]
+ prev_slots = mapping[prev_full]
+ self._register_c4_slots(mem_indexes[cont_rows], prev_slots + 1)
+
+ if new_rows:
+ self._realize_c4_pages(len(new_rows)) # 兑现: 精确需求, 复用已算的 new_rows
+ pages = self.mem_manager.alloc_c4_pages(len(new_rows)).to(mapping.device, non_blocking=True)
+ self._register_c4_slots(mem_indexes[new_rows], pages * page)
+ return
+
+ def _scatter_c128_slots(self, full_slots: torch.Tensor) -> None:
+ """为本批新组末 full 槽分配 c128 槽并写入映射。"""
+ if full_slots.numel() == 0:
+ return
+ full_slots = full_slots.reshape(-1)
+ self._realize_c128_slots(full_slots.numel())
+ new_slots = self.mem_manager.alloc_c128(full_slots.numel()).cuda(non_blocking=True)
+ self.mem_manager.full_to_c128_indexs[full_slots] = new_slots
+ return
+
+ def _realize_c4_pages(self, need_pages: int) -> None:
+ """压缩池兑现 —— 和主池在 prep 里调 free_radix_cache_to_get_enough_token 同一套路:
+ base_backend admission 已按"空闲+可回收"放行本步请求,这里在真分配前(scatter 已算好 need)
+ 把可回收的无引用 radix 节点驱逐出来腾出 c4 页,避免 alloc_c4_pages 触底 assert。
+ 可回收仍不足时由 admission 的 wait_pause 兜底。"""
+ if self.n_c4 == 0 or need_pages <= 0:
+ return
+ # 延迟 import: infer_batch 在模块顶 import 了 req_manager,顶层 import 会循环引用
+ from lightllm.server.router.model_infer.infer_batch import g_infer_context
+
+ if g_infer_context.radix_cache is not None:
+ g_infer_context.radix_cache.free_radix_cache_to_get_enough_c4_pages(need_pages)
+ return
+
+ def _realize_c128_slots(self, need_slots: int) -> None:
+ if self.n_c128 == 0 or need_slots <= 0:
+ return
+ from lightllm.server.router.model_infer.infer_batch import g_infer_context
+
+ if g_infer_context.radix_cache is not None:
+ g_infer_context.radix_cache.free_radix_cache_to_get_enough_c128_slots(need_slots)
+ return
+
+ def prepare_prefill_compress_slots(
+ self,
+ req_list: List[int],
+ ready_list: List[int],
+ seq_list: List[int],
+ mem_indexes: torch.Tensor,
+ ) -> None:
+ """prefill prep: 为本 chunk 内的组末 token(位置 (g+1)*ratio-1 ∈ [ready, seq))分配压缩槽，
+ 组末 full 槽直接从 generic preprocess 的 mem_indexes 取。"""
+ if self.n_c4 == 0 and self.n_c128 == 0:
+ return
+ if self.n_c4 > 0:
+ self._scatter_c4_prefill_slots_batched(req_list, ready_list, seq_list, mem_indexes)
+
+ if self.n_c128 > 0:
+ ratio = 128
+ full_offsets = []
+ mem_offset = 0
+ for req_idx, ready_len, seq_len in zip(req_list, ready_list, seq_list):
+ q_len = seq_len - ready_len
+ if req_idx == self.HOLD_REQUEST_ID:
+ mem_offset += q_len
+ continue
+ first, last = ready_len // ratio, seq_len // ratio
+ if last > first:
+ full_offsets.extend(
+ mem_offset + (entry + 1) * ratio - 1 - ready_len for entry in range(first, last)
+ )
+ mem_offset += q_len
+ if full_offsets:
+ offsets = g_pin_mem_manager.gen_from_list(
+ key="dsv4_c128_prefill_offsets", data=full_offsets, dtype=torch.int64
+ ).to(mem_indexes.device, non_blocking=True)
+ self._scatter_c128_slots(mem_indexes.reshape(-1)[offsets])
+ return
+
+ def prepare_decode_compress_slots(
+ self,
+ req_list: List[int],
+ seq_list: List[int],
+ mem_indexes: torch.Tensor,
+ prev_group_end_mem_indexes: Optional[torch.Tensor] = None,
+ ) -> None:
+ """decode prep: 本步 token 关闭一个组(seq_len % ratio == 0)时为其分配压缩槽并 scatter。
+ 组末 full 槽即本步的 mem_index。
+ 从 CPU 镜像读 seq_len/req_idx(host 算术,无 D2H);非关组步 rows 为空 => 不调 _scatter,零同步。"""
+ if self.n_c4 == 0 and self.n_c128 == 0:
+ return
+ if self.n_c4 > 0:
+ self._scatter_c4_decode_slots(
+ req_list,
+ seq_list,
+ mem_indexes,
+ prev_group_end_mem_indexes=prev_group_end_mem_indexes,
+ )
+
+ if self.n_c128 > 0:
+ ratio = 128
+ rows = [
+ i
+ for i, (req_idx, seq_len) in enumerate(zip(req_list, seq_list))
+ if req_idx != self.HOLD_REQUEST_ID and seq_len > 0 and seq_len % ratio == 0
+ ]
+ if rows:
+ self._scatter_c128_slots(mem_indexes.reshape(-1)[rows])
+ return
+
+ def alloc(self):
+ req_idx = super().alloc()
+ if req_idx is not None:
+ self.init_compress_state(req_idx)
+ return req_idx
+
+ def clear_runtime_state(self, req_idx: int):
+ # swa 槽位本身由 mem_manager.free 级联回收(随 full 槽位)，这里只复位出窗水位线。
+ self._swa_evict_marks[req_idx] = -1
+ return
+
+ def get_prompt_cache_value_ops(self):
+ return DeepseekV4PromptCacheValueOps(self)
+
+ def get_prompt_cache_page_size(self):
+ return DSV4_PROMPT_CACHE_PAGE_SIZE
+
+ def compute_swa_page_valid(self, full_slots: torch.Tensor) -> torch.Tensor:
+ """按当下 full_to_swa 映射给出按页有效性: full_slots [L](L 为 page 整数倍) ->
+ cpu bool [L/page]，页内全部映射有效才为 True。GPU gather + 同步,测试/校验用;
+ 插入热路径用 swa_page_valid_from_watermark(纯 CPU,免同步)。"""
+ page = self.get_prompt_cache_page_size()
+ assert full_slots.numel() % page == 0
+ if full_slots.numel() == 0:
+ return torch.zeros((0,), dtype=torch.bool)
+ swa = self.mem_manager.full_to_swa_indexs[full_slots.cuda().long().reshape(-1)]
+ return (swa.view(-1, page) >= 0).all(dim=1).cpu()
+
+ def swa_page_valid_from_watermark(self, req_idx: int, cache_len: int) -> torch.Tensor:
+ """插入时的按页有效性,纯 CPU: 请求自有 token 的 swa 映射只被出窗水位线回收
+ (阀不触活跃请求,级联只在 free 时),页 p 全驻留 ⟺ 页起点 page*p >= 水位线。
+
+ 与 compute_swa_page_valid 在插入时刻对自有 token 等价,但不做 GPU gather/同步——
+ router 关键路径上每次插入省一次对全部在途 kernel 的等待。bitmap 中借入前缀
+ ([0, ready) 的页)的行在 radix insert 切片时被丢弃(既有节点保留自己的 bitmap),
+ 其取值无影响。"""
+ page = self.get_prompt_cache_page_size()
+ mark = max(0, self._swa_evict_marks[req_idx])
+ n_pages = int(cache_len) // page
+ return torch.arange(n_pages, dtype=torch.long) * page >= mark
+
+ def slice_prompt_cache_payload(self, payload: DeepseekV4PromptCachePayload, start: int, end: int):
+ start = int(start)
+ end = int(end)
+ page = self.get_prompt_cache_page_size()
+ # radix page 保证分裂点页对齐，bitmap 可整页切分。
+ ans = DeepseekV4PromptCachePayload(
+ cache_len=end - start,
+ swa_page_valid=payload.swa_page_valid[start // page : end // page].clone()
+ if payload.swa_page_valid is not None
+ else None,
+ )
+ ans.refresh_swa_last_valid_page()
+ return ans
+
+ def concat_prompt_cache_payloads(self, payloads: List[DeepseekV4PromptCachePayload]):
+ if len(payloads) == 0:
+ return None
+ bitmaps = [p.swa_page_valid for p in payloads]
+ ans = DeepseekV4PromptCachePayload(
+ cache_len=sum(p.cache_len for p in payloads),
+ swa_page_valid=torch.cat(bitmaps, dim=0) if all(b is not None for b in bitmaps) else None,
+ )
+ if ans.swa_page_valid is None:
+ return ans
+
+ page = self.get_prompt_cache_page_size()
+ page_offset = 0
+ last_valid_page = -1
+ for item in payloads:
+ item_last = int(getattr(item, "swa_last_valid_page", -1))
+ if item_last >= 0:
+ last_valid_page = page_offset + item_last
+ page_offset += int(item.cache_len) // page
+ ans.swa_last_valid_page = last_valid_page
+ return ans
+
+ def build_prompt_cache_payload(
+ self,
+ cache_len: int,
+ ) -> DeepseekV4PromptCachePayload:
+ """构造插入载荷。compressor 状态不进载荷(c4 随 swa 页生灭、c128 边界自然归零),
+ cache_len 不再受序列末端约束——任意 128 对齐前缀皆可插入。
+ swa_page_valid 不在此填: 它必须用插入时刻的映射(infer batch 在 insert 前补)。"""
+ assert self.mem_manager is not None
+ return DeepseekV4PromptCachePayload(cache_len=int(cache_len))
+
+ def free(self, free_req_indexes, free_token_index):
+ """dense/swa/压缩槽全部经 mem_manager.free(free_token_index) 级联回收。"""
+ for req_index in free_req_indexes:
+ self.clear_runtime_state(req_index)
+ super().free(free_req_indexes, free_token_index)
+ return
+
+ def free_req(self, free_req_index: int):
+ self.clear_runtime_state(free_req_index)
+ return super().free_req(free_req_index)
+
+ def free_all(self):
+ super().free_all()
+ self._swa_evict_marks = [-1 for _ in range(self.max_request_num + 1)]
+ return
diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H100_80GB_HBM3/grouped_matmul:v1/{K=256,N=4096,expert_num=256,mul_routed_weight=true,out_dtype=torch.bfloat16,topk_num=1,use_fp8_w8a8=true}_NVIDIA_H100_80GB_HBM3.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H100_80GB_HBM3/grouped_matmul:v1/{K=256,N=4096,expert_num=256,mul_routed_weight=true,out_dtype=torch.bfloat16,topk_num=1,use_fp8_w8a8=true}_NVIDIA_H100_80GB_HBM3.json
new file mode 100644
index 0000000000..5204097669
--- /dev/null
+++ b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H100_80GB_HBM3/grouped_matmul:v1/{K=256,N=4096,expert_num=256,mul_routed_weight=true,out_dtype=torch.bfloat16,topk_num=1,use_fp8_w8a8=true}_NVIDIA_H100_80GB_HBM3.json
@@ -0,0 +1,110 @@
+{
+ "12288": {
+ "BLOCK_SIZE_K": 128,
+ "BLOCK_SIZE_M": 64,
+ "BLOCK_SIZE_N": 128,
+ "GROUP_SIZE_M": 64,
+ "NEED_TRANS": false,
+ "num_stages": 3,
+ "num_warps": 4
+ },
+ "1536": {
+ "BLOCK_SIZE_K": 64,
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 128,
+ "GROUP_SIZE_M": 1,
+ "NEED_TRANS": true,
+ "num_stages": 3,
+ "num_warps": 4
+ },
+ "192": {
+ "BLOCK_SIZE_K": 128,
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 128,
+ "GROUP_SIZE_M": 16,
+ "NEED_TRANS": true,
+ "num_stages": 2,
+ "num_warps": 4
+ },
+ "24576": {
+ "BLOCK_SIZE_K": 128,
+ "BLOCK_SIZE_M": 64,
+ "BLOCK_SIZE_N": 128,
+ "GROUP_SIZE_M": 32,
+ "NEED_TRANS": false,
+ "num_stages": 3,
+ "num_warps": 4
+ },
+ "384": {
+ "BLOCK_SIZE_K": 128,
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 128,
+ "GROUP_SIZE_M": 16,
+ "NEED_TRANS": true,
+ "num_stages": 2,
+ "num_warps": 4
+ },
+ "48": {
+ "BLOCK_SIZE_K": 64,
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 128,
+ "GROUP_SIZE_M": 64,
+ "NEED_TRANS": true,
+ "num_stages": 3,
+ "num_warps": 4
+ },
+ "49152": {
+ "BLOCK_SIZE_K": 128,
+ "BLOCK_SIZE_M": 64,
+ "BLOCK_SIZE_N": 128,
+ "GROUP_SIZE_M": 16,
+ "NEED_TRANS": false,
+ "num_stages": 3,
+ "num_warps": 4
+ },
+ "6": {
+ "BLOCK_SIZE_K": 64,
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 128,
+ "GROUP_SIZE_M": 32,
+ "NEED_TRANS": true,
+ "num_stages": 2,
+ "num_warps": 4
+ },
+ "600": {
+ "BLOCK_SIZE_K": 128,
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 128,
+ "GROUP_SIZE_M": 16,
+ "NEED_TRANS": true,
+ "num_stages": 2,
+ "num_warps": 4
+ },
+ "6144": {
+ "BLOCK_SIZE_K": 128,
+ "BLOCK_SIZE_M": 64,
+ "BLOCK_SIZE_N": 128,
+ "GROUP_SIZE_M": 64,
+ "NEED_TRANS": false,
+ "num_stages": 3,
+ "num_warps": 4
+ },
+ "768": {
+ "BLOCK_SIZE_K": 128,
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 128,
+ "GROUP_SIZE_M": 16,
+ "NEED_TRANS": true,
+ "num_stages": 2,
+ "num_warps": 4
+ },
+ "96": {
+ "BLOCK_SIZE_K": 128,
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 128,
+ "GROUP_SIZE_M": 64,
+ "NEED_TRANS": true,
+ "num_stages": 2,
+ "num_warps": 4
+ }
+}
\ No newline at end of file
diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H100_80GB_HBM3/grouped_matmul:v1/{K=4096,N=512,expert_num=256,mul_routed_weight=false,out_dtype=torch.bfloat16,topk_num=6,use_fp8_w8a8=true}_NVIDIA_H100_80GB_HBM3.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H100_80GB_HBM3/grouped_matmul:v1/{K=4096,N=512,expert_num=256,mul_routed_weight=false,out_dtype=torch.bfloat16,topk_num=6,use_fp8_w8a8=true}_NVIDIA_H100_80GB_HBM3.json
new file mode 100644
index 0000000000..ac4ce1ba57
--- /dev/null
+++ b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H100_80GB_HBM3/grouped_matmul:v1/{K=4096,N=512,expert_num=256,mul_routed_weight=false,out_dtype=torch.bfloat16,topk_num=6,use_fp8_w8a8=true}_NVIDIA_H100_80GB_HBM3.json
@@ -0,0 +1,110 @@
+{
+ "1": {
+ "BLOCK_SIZE_K": 128,
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 64,
+ "GROUP_SIZE_M": 1,
+ "NEED_TRANS": true,
+ "num_stages": 4,
+ "num_warps": 4
+ },
+ "100": {
+ "BLOCK_SIZE_K": 128,
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 64,
+ "GROUP_SIZE_M": 1,
+ "NEED_TRANS": true,
+ "num_stages": 3,
+ "num_warps": 4
+ },
+ "1024": {
+ "BLOCK_SIZE_K": 128,
+ "BLOCK_SIZE_M": 64,
+ "BLOCK_SIZE_N": 64,
+ "GROUP_SIZE_M": 1,
+ "NEED_TRANS": false,
+ "num_stages": 3,
+ "num_warps": 4
+ },
+ "128": {
+ "BLOCK_SIZE_K": 128,
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 64,
+ "GROUP_SIZE_M": 1,
+ "NEED_TRANS": true,
+ "num_stages": 5,
+ "num_warps": 4
+ },
+ "16": {
+ "BLOCK_SIZE_K": 128,
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 64,
+ "GROUP_SIZE_M": 1,
+ "NEED_TRANS": true,
+ "num_stages": 4,
+ "num_warps": 4
+ },
+ "2048": {
+ "BLOCK_SIZE_K": 128,
+ "BLOCK_SIZE_M": 64,
+ "BLOCK_SIZE_N": 128,
+ "GROUP_SIZE_M": 1,
+ "NEED_TRANS": false,
+ "num_stages": 3,
+ "num_warps": 8
+ },
+ "256": {
+ "BLOCK_SIZE_K": 128,
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 128,
+ "GROUP_SIZE_M": 1,
+ "NEED_TRANS": true,
+ "num_stages": 3,
+ "num_warps": 4
+ },
+ "32": {
+ "BLOCK_SIZE_K": 128,
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 64,
+ "GROUP_SIZE_M": 16,
+ "NEED_TRANS": true,
+ "num_stages": 5,
+ "num_warps": 4
+ },
+ "4096": {
+ "BLOCK_SIZE_K": 128,
+ "BLOCK_SIZE_M": 128,
+ "BLOCK_SIZE_N": 128,
+ "GROUP_SIZE_M": 1,
+ "NEED_TRANS": false,
+ "num_stages": 5,
+ "num_warps": 8
+ },
+ "64": {
+ "BLOCK_SIZE_K": 128,
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 64,
+ "GROUP_SIZE_M": 1,
+ "NEED_TRANS": true,
+ "num_stages": 5,
+ "num_warps": 4
+ },
+ "8": {
+ "BLOCK_SIZE_K": 128,
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 64,
+ "GROUP_SIZE_M": 32,
+ "NEED_TRANS": true,
+ "num_stages": 4,
+ "num_warps": 4
+ },
+ "8192": {
+ "BLOCK_SIZE_K": 128,
+ "BLOCK_SIZE_M": 64,
+ "BLOCK_SIZE_N": 128,
+ "GROUP_SIZE_M": 1,
+ "NEED_TRANS": false,
+ "num_stages": 3,
+ "num_warps": 4
+ }
+}
\ No newline at end of file
diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H100_80GB_HBM3/moe_align_fused:v1/{topk_num=6}_NVIDIA_H100_80GB_HBM3.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H100_80GB_HBM3/moe_align_fused:v1/{topk_num=6}_NVIDIA_H100_80GB_HBM3.json
new file mode 100644
index 0000000000..6aa8d18c54
--- /dev/null
+++ b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H100_80GB_HBM3/moe_align_fused:v1/{topk_num=6}_NVIDIA_H100_80GB_HBM3.json
@@ -0,0 +1,50 @@
+{
+ "1": {
+ "BLOCK_SIZE": 256,
+ "num_warps": 2
+ },
+ "100": {
+ "BLOCK_SIZE": 128,
+ "num_warps": 8
+ },
+ "1024": {
+ "BLOCK_SIZE": 128,
+ "num_warps": 8
+ },
+ "128": {
+ "BLOCK_SIZE": 128,
+ "num_warps": 8
+ },
+ "16": {
+ "BLOCK_SIZE": 512,
+ "num_warps": 4
+ },
+ "2048": {
+ "BLOCK_SIZE": 128,
+ "num_warps": 4
+ },
+ "256": {
+ "BLOCK_SIZE": 128,
+ "num_warps": 8
+ },
+ "32": {
+ "BLOCK_SIZE": 128,
+ "num_warps": 8
+ },
+ "4096": {
+ "BLOCK_SIZE": 256,
+ "num_warps": 4
+ },
+ "64": {
+ "BLOCK_SIZE": 128,
+ "num_warps": 8
+ },
+ "8": {
+ "BLOCK_SIZE": 256,
+ "num_warps": 2
+ },
+ "8192": {
+ "BLOCK_SIZE": 256,
+ "num_warps": 4
+ }
+}
\ No newline at end of file
diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H100_80GB_HBM3/moe_sum_reduce:v1/{hidden_dim=4096,out_dtype=torch.bfloat16,topk_num=6}_NVIDIA_H100_80GB_HBM3.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H100_80GB_HBM3/moe_sum_reduce:v1/{hidden_dim=4096,out_dtype=torch.bfloat16,topk_num=6}_NVIDIA_H100_80GB_HBM3.json
new file mode 100644
index 0000000000..e2da8bc968
--- /dev/null
+++ b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H100_80GB_HBM3/moe_sum_reduce:v1/{hidden_dim=4096,out_dtype=torch.bfloat16,topk_num=6}_NVIDIA_H100_80GB_HBM3.json
@@ -0,0 +1,74 @@
+{
+ "1": {
+ "BLOCK_DIM": 256,
+ "BLOCK_M": 1,
+ "NUM_STAGE": 1,
+ "num_warps": 8
+ },
+ "100": {
+ "BLOCK_DIM": 1024,
+ "BLOCK_M": 1,
+ "NUM_STAGE": 1,
+ "num_warps": 4
+ },
+ "1024": {
+ "BLOCK_DIM": 256,
+ "BLOCK_M": 1,
+ "NUM_STAGE": 4,
+ "num_warps": 1
+ },
+ "128": {
+ "BLOCK_DIM": 512,
+ "BLOCK_M": 1,
+ "NUM_STAGE": 1,
+ "num_warps": 4
+ },
+ "16": {
+ "BLOCK_DIM": 512,
+ "BLOCK_M": 1,
+ "NUM_STAGE": 1,
+ "num_warps": 4
+ },
+ "2048": {
+ "BLOCK_DIM": 1024,
+ "BLOCK_M": 1,
+ "NUM_STAGE": 1,
+ "num_warps": 8
+ },
+ "256": {
+ "BLOCK_DIM": 512,
+ "BLOCK_M": 1,
+ "NUM_STAGE": 1,
+ "num_warps": 4
+ },
+ "32": {
+ "BLOCK_DIM": 512,
+ "BLOCK_M": 1,
+ "NUM_STAGE": 1,
+ "num_warps": 4
+ },
+ "4096": {
+ "BLOCK_DIM": 512,
+ "BLOCK_M": 1,
+ "NUM_STAGE": 4,
+ "num_warps": 2
+ },
+ "64": {
+ "BLOCK_DIM": 512,
+ "BLOCK_M": 1,
+ "NUM_STAGE": 1,
+ "num_warps": 4
+ },
+ "8": {
+ "BLOCK_DIM": 64,
+ "BLOCK_M": 1,
+ "NUM_STAGE": 1,
+ "num_warps": 2
+ },
+ "8192": {
+ "BLOCK_DIM": 1024,
+ "BLOCK_M": 1,
+ "NUM_STAGE": 1,
+ "num_warps": 8
+ }
+}
\ No newline at end of file
diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H100_80GB_HBM3/rotary_emb_fwd:v1/{HEAD_DIM=64,K_HEAD_NUM=0,Q_HEAD_NUM=64,dtype=torch.bfloat16}_NVIDIA_H100_80GB_HBM3.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H100_80GB_HBM3/rotary_emb_fwd:v1/{HEAD_DIM=64,K_HEAD_NUM=0,Q_HEAD_NUM=64,dtype=torch.bfloat16}_NVIDIA_H100_80GB_HBM3.json
new file mode 100644
index 0000000000..e700378de1
--- /dev/null
+++ b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H100_80GB_HBM3/rotary_emb_fwd:v1/{HEAD_DIM=64,K_HEAD_NUM=0,Q_HEAD_NUM=64,dtype=torch.bfloat16}_NVIDIA_H100_80GB_HBM3.json
@@ -0,0 +1,74 @@
+{
+ "1": {
+ "BLOCK_SEQ": 8,
+ "HEAD_PARALLEL_NUM": 16,
+ "num_stages": 4,
+ "num_warps": 4
+ },
+ "100": {
+ "BLOCK_SEQ": 1,
+ "HEAD_PARALLEL_NUM": 16,
+ "num_stages": 2,
+ "num_warps": 2
+ },
+ "1024": {
+ "BLOCK_SEQ": 1,
+ "HEAD_PARALLEL_NUM": 4,
+ "num_stages": 5,
+ "num_warps": 1
+ },
+ "128": {
+ "BLOCK_SEQ": 1,
+ "HEAD_PARALLEL_NUM": 16,
+ "num_stages": 2,
+ "num_warps": 2
+ },
+ "16": {
+ "BLOCK_SEQ": 1,
+ "HEAD_PARALLEL_NUM": 16,
+ "num_stages": 2,
+ "num_warps": 8
+ },
+ "2048": {
+ "BLOCK_SEQ": 1,
+ "HEAD_PARALLEL_NUM": 4,
+ "num_stages": 3,
+ "num_warps": 1
+ },
+ "256": {
+ "BLOCK_SEQ": 1,
+ "HEAD_PARALLEL_NUM": 8,
+ "num_stages": 2,
+ "num_warps": 2
+ },
+ "32": {
+ "BLOCK_SEQ": 1,
+ "HEAD_PARALLEL_NUM": 16,
+ "num_stages": 2,
+ "num_warps": 2
+ },
+ "4096": {
+ "BLOCK_SEQ": 2,
+ "HEAD_PARALLEL_NUM": 2,
+ "num_stages": 4,
+ "num_warps": 1
+ },
+ "64": {
+ "BLOCK_SEQ": 1,
+ "HEAD_PARALLEL_NUM": 16,
+ "num_stages": 2,
+ "num_warps": 2
+ },
+ "8": {
+ "BLOCK_SEQ": 1,
+ "HEAD_PARALLEL_NUM": 16,
+ "num_stages": 2,
+ "num_warps": 8
+ },
+ "8192": {
+ "BLOCK_SEQ": 8,
+ "HEAD_PARALLEL_NUM": 4,
+ "num_stages": 4,
+ "num_warps": 1
+ }
+}
\ No newline at end of file
diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H100_80GB_HBM3/rotary_emb_fwd:v1/{HEAD_DIM=64,K_HEAD_NUM=0,Q_HEAD_NUM=8,dtype=torch.bfloat16}_NVIDIA_H100_80GB_HBM3.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H100_80GB_HBM3/rotary_emb_fwd:v1/{HEAD_DIM=64,K_HEAD_NUM=0,Q_HEAD_NUM=8,dtype=torch.bfloat16}_NVIDIA_H100_80GB_HBM3.json
new file mode 100644
index 0000000000..588fd4a934
--- /dev/null
+++ b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H100_80GB_HBM3/rotary_emb_fwd:v1/{HEAD_DIM=64,K_HEAD_NUM=0,Q_HEAD_NUM=8,dtype=torch.bfloat16}_NVIDIA_H100_80GB_HBM3.json
@@ -0,0 +1,74 @@
+{
+ "1": {
+ "BLOCK_SEQ": 1,
+ "HEAD_PARALLEL_NUM": 8,
+ "num_stages": 5,
+ "num_warps": 8
+ },
+ "100": {
+ "BLOCK_SEQ": 1,
+ "HEAD_PARALLEL_NUM": 8,
+ "num_stages": 1,
+ "num_warps": 4
+ },
+ "1024": {
+ "BLOCK_SEQ": 1,
+ "HEAD_PARALLEL_NUM": 1,
+ "num_stages": 3,
+ "num_warps": 4
+ },
+ "128": {
+ "BLOCK_SEQ": 1,
+ "HEAD_PARALLEL_NUM": 8,
+ "num_stages": 1,
+ "num_warps": 4
+ },
+ "16": {
+ "BLOCK_SEQ": 1,
+ "HEAD_PARALLEL_NUM": 8,
+ "num_stages": 4,
+ "num_warps": 8
+ },
+ "2048": {
+ "BLOCK_SEQ": 1,
+ "HEAD_PARALLEL_NUM": 1,
+ "num_stages": 4,
+ "num_warps": 2
+ },
+ "256": {
+ "BLOCK_SEQ": 2,
+ "HEAD_PARALLEL_NUM": 8,
+ "num_stages": 5,
+ "num_warps": 1
+ },
+ "32": {
+ "BLOCK_SEQ": 1,
+ "HEAD_PARALLEL_NUM": 8,
+ "num_stages": 2,
+ "num_warps": 8
+ },
+ "4096": {
+ "BLOCK_SEQ": 1,
+ "HEAD_PARALLEL_NUM": 1,
+ "num_stages": 1,
+ "num_warps": 2
+ },
+ "64": {
+ "BLOCK_SEQ": 1,
+ "HEAD_PARALLEL_NUM": 16,
+ "num_stages": 2,
+ "num_warps": 4
+ },
+ "8": {
+ "BLOCK_SEQ": 1,
+ "HEAD_PARALLEL_NUM": 8,
+ "num_stages": 5,
+ "num_warps": 8
+ },
+ "8192": {
+ "BLOCK_SEQ": 2,
+ "HEAD_PARALLEL_NUM": 1,
+ "num_stages": 3,
+ "num_warps": 1
+ }
+}
\ No newline at end of file
diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H100_80GB_HBM3/silu_and_mul_fwd:v1/{N=2048,out_dtype=torch.bfloat16}_NVIDIA_H100_80GB_HBM3.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H100_80GB_HBM3/silu_and_mul_fwd:v1/{N=2048,out_dtype=torch.bfloat16}_NVIDIA_H100_80GB_HBM3.json
new file mode 100644
index 0000000000..fbd605f99a
--- /dev/null
+++ b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H100_80GB_HBM3/silu_and_mul_fwd:v1/{N=2048,out_dtype=torch.bfloat16}_NVIDIA_H100_80GB_HBM3.json
@@ -0,0 +1,632 @@
+{
+ "1": {
+ "BLOCK_M": 64,
+ "BLOCK_N": 32,
+ "NUM_STAGES": 2,
+ "num_warps": 8
+ },
+ "100": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 2,
+ "num_warps": 4
+ },
+ "1024": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "1152": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 1,
+ "num_warps": 1
+ },
+ "12160": {
+ "BLOCK_M": 64,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "128": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 1,
+ "num_warps": 1
+ },
+ "1280": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 1,
+ "num_warps": 1
+ },
+ "12928": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 2,
+ "num_warps": 1
+ },
+ "13184": {
+ "BLOCK_M": 32,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "13312": {
+ "BLOCK_M": 64,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "13568": {
+ "BLOCK_M": 64,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "13824": {
+ "BLOCK_M": 64,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "13952": {
+ "BLOCK_M": 64,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "1408": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "14080": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 1,
+ "num_warps": 1
+ },
+ "14336": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 2,
+ "num_warps": 1
+ },
+ "14592": {
+ "BLOCK_M": 32,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 1,
+ "num_warps": 1
+ },
+ "14720": {
+ "BLOCK_M": 32,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 1,
+ "num_warps": 1
+ },
+ "14848": {
+ "BLOCK_M": 64,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "14976": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 1,
+ "num_warps": 1
+ },
+ "15232": {
+ "BLOCK_M": 32,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 1,
+ "num_warps": 1
+ },
+ "1536": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "16": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 2,
+ "num_warps": 4
+ },
+ "1664": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "1792": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "1920": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "2048": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "2176": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "2304": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "24192": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "2432": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "24960": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "25088": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "25344": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 2,
+ "num_warps": 1
+ },
+ "25472": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "256": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 1,
+ "num_warps": 1
+ },
+ "2560": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "25728": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "25856": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "26112": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "26240": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "26752": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "2688": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "27008": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "27136": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "27520": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "27776": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "27904": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "28032": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "2816": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "28288": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "28800": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "2944": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "3072": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "32": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 4
+ },
+ "3200": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "3328": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "3456": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "3584": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "3712": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "384": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "3840": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 2,
+ "num_warps": 1
+ },
+ "3968": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "4096": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "4224": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 1,
+ "num_warps": 1
+ },
+ "4352": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "4480": {
+ "BLOCK_M": 32,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "46336": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "46720": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "46976": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "48896": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "49152": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "49408": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "50048": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "50304": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "50560": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "50688": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "50816": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "50944": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "512": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "51328": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "52608": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "53248": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "53632": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "54144": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "54272": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "54528": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "54656": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "55040": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "64": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 2,
+ "num_warps": 1
+ },
+ "640": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "7424": {
+ "BLOCK_M": 32,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "7552": {
+ "BLOCK_M": 32,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "768": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "7680": {
+ "BLOCK_M": 32,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "7808": {
+ "BLOCK_M": 32,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "7936": {
+ "BLOCK_M": 32,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "8": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 2,
+ "num_warps": 4
+ },
+ "8064": {
+ "BLOCK_M": 32,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "8192": {
+ "BLOCK_M": 32,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "8320": {
+ "BLOCK_M": 32,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "8448": {
+ "BLOCK_M": 32,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "8576": {
+ "BLOCK_M": 32,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "896": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 1,
+ "num_warps": 4
+ },
+ "8960": {
+ "BLOCK_M": 32,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ }
+}
\ No newline at end of file
diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H100_80GB_HBM3/silu_and_mul_fwd:v1/{N=256,out_dtype=torch.bfloat16}_NVIDIA_H100_80GB_HBM3.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H100_80GB_HBM3/silu_and_mul_fwd:v1/{N=256,out_dtype=torch.bfloat16}_NVIDIA_H100_80GB_HBM3.json
new file mode 100644
index 0000000000..4d7e8f1183
--- /dev/null
+++ b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H100_80GB_HBM3/silu_and_mul_fwd:v1/{N=256,out_dtype=torch.bfloat16}_NVIDIA_H100_80GB_HBM3.json
@@ -0,0 +1,146 @@
+{
+ "1": {
+ "BLOCK_M": 128,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 2,
+ "num_warps": 4
+ },
+ "100": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 1,
+ "num_warps": 1
+ },
+ "1024": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 1,
+ "num_warps": 4
+ },
+ "12288": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "128": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 2,
+ "num_warps": 4
+ },
+ "1536": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 2,
+ "num_warps": 1
+ },
+ "16": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 4
+ },
+ "192": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 1,
+ "num_warps": 4
+ },
+ "2048": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 1,
+ "num_warps": 1
+ },
+ "24576": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "256": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "32": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 32,
+ "NUM_STAGES": 2,
+ "num_warps": 4
+ },
+ "384": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "4096": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "48": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 2,
+ "num_warps": 1
+ },
+ "49152": {
+ "BLOCK_M": 32,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "6": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 8
+ },
+ "600": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "6144": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "64": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 8
+ },
+ "768": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 1,
+ "num_warps": 4
+ },
+ "8": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 32,
+ "NUM_STAGES": 1,
+ "num_warps": 8
+ },
+ "8192": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "96": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 8
+ }
+}
\ No newline at end of file
diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H200/grouped_matmul:v1/{K=256,N=4096,expert_num=256,mul_routed_weight=true,out_dtype=torch.bfloat16,topk_num=1,use_fp8_w8a8=true}_NVIDIA_H200.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H200/grouped_matmul:v1/{K=256,N=4096,expert_num=256,mul_routed_weight=true,out_dtype=torch.bfloat16,topk_num=1,use_fp8_w8a8=true}_NVIDIA_H200.json
new file mode 100644
index 0000000000..b1aae6bfba
--- /dev/null
+++ b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H200/grouped_matmul:v1/{K=256,N=4096,expert_num=256,mul_routed_weight=true,out_dtype=torch.bfloat16,topk_num=1,use_fp8_w8a8=true}_NVIDIA_H200.json
@@ -0,0 +1,110 @@
+{
+ "12288": {
+ "BLOCK_SIZE_K": 128,
+ "BLOCK_SIZE_M": 64,
+ "BLOCK_SIZE_N": 128,
+ "GROUP_SIZE_M": 64,
+ "NEED_TRANS": false,
+ "num_stages": 3,
+ "num_warps": 4
+ },
+ "1536": {
+ "BLOCK_SIZE_K": 128,
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 128,
+ "GROUP_SIZE_M": 64,
+ "NEED_TRANS": true,
+ "num_stages": 2,
+ "num_warps": 4
+ },
+ "192": {
+ "BLOCK_SIZE_K": 64,
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 128,
+ "GROUP_SIZE_M": 32,
+ "NEED_TRANS": true,
+ "num_stages": 3,
+ "num_warps": 4
+ },
+ "24576": {
+ "BLOCK_SIZE_K": 128,
+ "BLOCK_SIZE_M": 64,
+ "BLOCK_SIZE_N": 128,
+ "GROUP_SIZE_M": 16,
+ "NEED_TRANS": false,
+ "num_stages": 3,
+ "num_warps": 4
+ },
+ "384": {
+ "BLOCK_SIZE_K": 64,
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 128,
+ "GROUP_SIZE_M": 32,
+ "NEED_TRANS": true,
+ "num_stages": 3,
+ "num_warps": 4
+ },
+ "48": {
+ "BLOCK_SIZE_K": 64,
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 128,
+ "GROUP_SIZE_M": 64,
+ "NEED_TRANS": true,
+ "num_stages": 3,
+ "num_warps": 4
+ },
+ "49152": {
+ "BLOCK_SIZE_K": 128,
+ "BLOCK_SIZE_M": 64,
+ "BLOCK_SIZE_N": 128,
+ "GROUP_SIZE_M": 16,
+ "NEED_TRANS": false,
+ "num_stages": 3,
+ "num_warps": 4
+ },
+ "6": {
+ "BLOCK_SIZE_K": 64,
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 128,
+ "GROUP_SIZE_M": 32,
+ "NEED_TRANS": true,
+ "num_stages": 3,
+ "num_warps": 4
+ },
+ "600": {
+ "BLOCK_SIZE_K": 64,
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 128,
+ "GROUP_SIZE_M": 64,
+ "NEED_TRANS": true,
+ "num_stages": 3,
+ "num_warps": 4
+ },
+ "6144": {
+ "BLOCK_SIZE_K": 128,
+ "BLOCK_SIZE_M": 32,
+ "BLOCK_SIZE_N": 128,
+ "GROUP_SIZE_M": 64,
+ "NEED_TRANS": true,
+ "num_stages": 2,
+ "num_warps": 4
+ },
+ "768": {
+ "BLOCK_SIZE_K": 64,
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 128,
+ "GROUP_SIZE_M": 32,
+ "NEED_TRANS": true,
+ "num_stages": 3,
+ "num_warps": 4
+ },
+ "96": {
+ "BLOCK_SIZE_K": 64,
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 128,
+ "GROUP_SIZE_M": 64,
+ "NEED_TRANS": true,
+ "num_stages": 3,
+ "num_warps": 4
+ }
+}
\ No newline at end of file
diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H200/grouped_matmul:v1/{K=4096,N=512,expert_num=256,mul_routed_weight=false,out_dtype=torch.bfloat16,topk_num=6,use_fp8_w8a8=true}_NVIDIA_H200.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H200/grouped_matmul:v1/{K=4096,N=512,expert_num=256,mul_routed_weight=false,out_dtype=torch.bfloat16,topk_num=6,use_fp8_w8a8=true}_NVIDIA_H200.json
new file mode 100644
index 0000000000..9ffb0efd19
--- /dev/null
+++ b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H200/grouped_matmul:v1/{K=4096,N=512,expert_num=256,mul_routed_weight=false,out_dtype=torch.bfloat16,topk_num=6,use_fp8_w8a8=true}_NVIDIA_H200.json
@@ -0,0 +1,110 @@
+{
+ "1": {
+ "BLOCK_SIZE_K": 128,
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 64,
+ "GROUP_SIZE_M": 1,
+ "NEED_TRANS": true,
+ "num_stages": 3,
+ "num_warps": 4
+ },
+ "100": {
+ "BLOCK_SIZE_K": 128,
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 64,
+ "GROUP_SIZE_M": 16,
+ "NEED_TRANS": true,
+ "num_stages": 3,
+ "num_warps": 4
+ },
+ "1024": {
+ "BLOCK_SIZE_K": 128,
+ "BLOCK_SIZE_M": 32,
+ "BLOCK_SIZE_N": 128,
+ "GROUP_SIZE_M": 16,
+ "NEED_TRANS": true,
+ "num_stages": 3,
+ "num_warps": 4
+ },
+ "128": {
+ "BLOCK_SIZE_K": 128,
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 128,
+ "GROUP_SIZE_M": 1,
+ "NEED_TRANS": true,
+ "num_stages": 5,
+ "num_warps": 4
+ },
+ "16": {
+ "BLOCK_SIZE_K": 128,
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 64,
+ "GROUP_SIZE_M": 32,
+ "NEED_TRANS": true,
+ "num_stages": 3,
+ "num_warps": 4
+ },
+ "2048": {
+ "BLOCK_SIZE_K": 128,
+ "BLOCK_SIZE_M": 64,
+ "BLOCK_SIZE_N": 128,
+ "GROUP_SIZE_M": 16,
+ "NEED_TRANS": false,
+ "num_stages": 3,
+ "num_warps": 4
+ },
+ "256": {
+ "BLOCK_SIZE_K": 128,
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 128,
+ "GROUP_SIZE_M": 1,
+ "NEED_TRANS": true,
+ "num_stages": 3,
+ "num_warps": 4
+ },
+ "32": {
+ "BLOCK_SIZE_K": 128,
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 64,
+ "GROUP_SIZE_M": 1,
+ "NEED_TRANS": true,
+ "num_stages": 4,
+ "num_warps": 4
+ },
+ "4096": {
+ "BLOCK_SIZE_K": 128,
+ "BLOCK_SIZE_M": 128,
+ "BLOCK_SIZE_N": 128,
+ "GROUP_SIZE_M": 16,
+ "NEED_TRANS": false,
+ "num_stages": 5,
+ "num_warps": 8
+ },
+ "64": {
+ "BLOCK_SIZE_K": 128,
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 64,
+ "GROUP_SIZE_M": 1,
+ "NEED_TRANS": true,
+ "num_stages": 5,
+ "num_warps": 4
+ },
+ "8": {
+ "BLOCK_SIZE_K": 128,
+ "BLOCK_SIZE_M": 16,
+ "BLOCK_SIZE_N": 64,
+ "GROUP_SIZE_M": 16,
+ "NEED_TRANS": true,
+ "num_stages": 4,
+ "num_warps": 4
+ },
+ "8192": {
+ "BLOCK_SIZE_K": 128,
+ "BLOCK_SIZE_M": 64,
+ "BLOCK_SIZE_N": 128,
+ "GROUP_SIZE_M": 1,
+ "NEED_TRANS": false,
+ "num_stages": 3,
+ "num_warps": 4
+ }
+}
\ No newline at end of file
diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H200/moe_align_fused:v1/{topk_num=6}_NVIDIA_H200.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H200/moe_align_fused:v1/{topk_num=6}_NVIDIA_H200.json
new file mode 100644
index 0000000000..85a20d9b1b
--- /dev/null
+++ b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H200/moe_align_fused:v1/{topk_num=6}_NVIDIA_H200.json
@@ -0,0 +1,50 @@
+{
+ "1": {
+ "BLOCK_SIZE": 128,
+ "num_warps": 1
+ },
+ "100": {
+ "BLOCK_SIZE": 128,
+ "num_warps": 4
+ },
+ "1024": {
+ "BLOCK_SIZE": 128,
+ "num_warps": 8
+ },
+ "128": {
+ "BLOCK_SIZE": 128,
+ "num_warps": 4
+ },
+ "16": {
+ "BLOCK_SIZE": 256,
+ "num_warps": 8
+ },
+ "2048": {
+ "BLOCK_SIZE": 128,
+ "num_warps": 8
+ },
+ "256": {
+ "BLOCK_SIZE": 128,
+ "num_warps": 8
+ },
+ "32": {
+ "BLOCK_SIZE": 128,
+ "num_warps": 8
+ },
+ "4096": {
+ "BLOCK_SIZE": 128,
+ "num_warps": 8
+ },
+ "64": {
+ "BLOCK_SIZE": 128,
+ "num_warps": 4
+ },
+ "8": {
+ "BLOCK_SIZE": 512,
+ "num_warps": 4
+ },
+ "8192": {
+ "BLOCK_SIZE": 256,
+ "num_warps": 8
+ }
+}
\ No newline at end of file
diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H200/moe_sum_reduce:v1/{hidden_dim=4096,out_dtype=torch.bfloat16,topk_num=6}_NVIDIA_H200.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H200/moe_sum_reduce:v1/{hidden_dim=4096,out_dtype=torch.bfloat16,topk_num=6}_NVIDIA_H200.json
new file mode 100644
index 0000000000..de2f015a04
--- /dev/null
+++ b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H200/moe_sum_reduce:v1/{hidden_dim=4096,out_dtype=torch.bfloat16,topk_num=6}_NVIDIA_H200.json
@@ -0,0 +1,74 @@
+{
+ "1": {
+ "BLOCK_DIM": 128,
+ "BLOCK_M": 1,
+ "NUM_STAGE": 2,
+ "num_warps": 4
+ },
+ "100": {
+ "BLOCK_DIM": 1024,
+ "BLOCK_M": 1,
+ "NUM_STAGE": 1,
+ "num_warps": 4
+ },
+ "1024": {
+ "BLOCK_DIM": 256,
+ "BLOCK_M": 4,
+ "NUM_STAGE": 4,
+ "num_warps": 1
+ },
+ "128": {
+ "BLOCK_DIM": 1024,
+ "BLOCK_M": 1,
+ "NUM_STAGE": 1,
+ "num_warps": 8
+ },
+ "16": {
+ "BLOCK_DIM": 1024,
+ "BLOCK_M": 1,
+ "NUM_STAGE": 1,
+ "num_warps": 8
+ },
+ "2048": {
+ "BLOCK_DIM": 512,
+ "BLOCK_M": 1,
+ "NUM_STAGE": 4,
+ "num_warps": 2
+ },
+ "256": {
+ "BLOCK_DIM": 1024,
+ "BLOCK_M": 1,
+ "NUM_STAGE": 2,
+ "num_warps": 1
+ },
+ "32": {
+ "BLOCK_DIM": 256,
+ "BLOCK_M": 1,
+ "NUM_STAGE": 1,
+ "num_warps": 2
+ },
+ "4096": {
+ "BLOCK_DIM": 1024,
+ "BLOCK_M": 1,
+ "NUM_STAGE": 4,
+ "num_warps": 1
+ },
+ "64": {
+ "BLOCK_DIM": 1024,
+ "BLOCK_M": 1,
+ "NUM_STAGE": 1,
+ "num_warps": 8
+ },
+ "8": {
+ "BLOCK_DIM": 512,
+ "BLOCK_M": 1,
+ "NUM_STAGE": 1,
+ "num_warps": 4
+ },
+ "8192": {
+ "BLOCK_DIM": 1024,
+ "BLOCK_M": 1,
+ "NUM_STAGE": 4,
+ "num_warps": 2
+ }
+}
\ No newline at end of file
diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H200/rotary_emb_fwd:v1/{HEAD_DIM=64,K_HEAD_NUM=0,Q_HEAD_NUM=64,dtype=torch.bfloat16}_NVIDIA_H200.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H200/rotary_emb_fwd:v1/{HEAD_DIM=64,K_HEAD_NUM=0,Q_HEAD_NUM=64,dtype=torch.bfloat16}_NVIDIA_H200.json
new file mode 100644
index 0000000000..e40e19975e
--- /dev/null
+++ b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H200/rotary_emb_fwd:v1/{HEAD_DIM=64,K_HEAD_NUM=0,Q_HEAD_NUM=64,dtype=torch.bfloat16}_NVIDIA_H200.json
@@ -0,0 +1,74 @@
+{
+ "1": {
+ "BLOCK_SEQ": 2,
+ "HEAD_PARALLEL_NUM": 16,
+ "num_stages": 4,
+ "num_warps": 8
+ },
+ "100": {
+ "BLOCK_SEQ": 1,
+ "HEAD_PARALLEL_NUM": 16,
+ "num_stages": 3,
+ "num_warps": 1
+ },
+ "1024": {
+ "BLOCK_SEQ": 1,
+ "HEAD_PARALLEL_NUM": 4,
+ "num_stages": 1,
+ "num_warps": 2
+ },
+ "128": {
+ "BLOCK_SEQ": 1,
+ "HEAD_PARALLEL_NUM": 16,
+ "num_stages": 3,
+ "num_warps": 1
+ },
+ "16": {
+ "BLOCK_SEQ": 1,
+ "HEAD_PARALLEL_NUM": 16,
+ "num_stages": 2,
+ "num_warps": 4
+ },
+ "2048": {
+ "BLOCK_SEQ": 1,
+ "HEAD_PARALLEL_NUM": 2,
+ "num_stages": 5,
+ "num_warps": 1
+ },
+ "256": {
+ "BLOCK_SEQ": 1,
+ "HEAD_PARALLEL_NUM": 8,
+ "num_stages": 5,
+ "num_warps": 2
+ },
+ "32": {
+ "BLOCK_SEQ": 1,
+ "HEAD_PARALLEL_NUM": 16,
+ "num_stages": 5,
+ "num_warps": 4
+ },
+ "4096": {
+ "BLOCK_SEQ": 1,
+ "HEAD_PARALLEL_NUM": 1,
+ "num_stages": 3,
+ "num_warps": 1
+ },
+ "64": {
+ "BLOCK_SEQ": 1,
+ "HEAD_PARALLEL_NUM": 16,
+ "num_stages": 1,
+ "num_warps": 1
+ },
+ "8": {
+ "BLOCK_SEQ": 1,
+ "HEAD_PARALLEL_NUM": 16,
+ "num_stages": 2,
+ "num_warps": 1
+ },
+ "8192": {
+ "BLOCK_SEQ": 2,
+ "HEAD_PARALLEL_NUM": 1,
+ "num_stages": 2,
+ "num_warps": 2
+ }
+}
\ No newline at end of file
diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H200/rotary_emb_fwd:v1/{HEAD_DIM=64,K_HEAD_NUM=0,Q_HEAD_NUM=8,dtype=torch.bfloat16}_NVIDIA_H200.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H200/rotary_emb_fwd:v1/{HEAD_DIM=64,K_HEAD_NUM=0,Q_HEAD_NUM=8,dtype=torch.bfloat16}_NVIDIA_H200.json
new file mode 100644
index 0000000000..88742a0b13
--- /dev/null
+++ b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H200/rotary_emb_fwd:v1/{HEAD_DIM=64,K_HEAD_NUM=0,Q_HEAD_NUM=8,dtype=torch.bfloat16}_NVIDIA_H200.json
@@ -0,0 +1,74 @@
+{
+ "1": {
+ "BLOCK_SEQ": 1,
+ "HEAD_PARALLEL_NUM": 16,
+ "num_stages": 4,
+ "num_warps": 1
+ },
+ "100": {
+ "BLOCK_SEQ": 2,
+ "HEAD_PARALLEL_NUM": 8,
+ "num_stages": 5,
+ "num_warps": 2
+ },
+ "1024": {
+ "BLOCK_SEQ": 1,
+ "HEAD_PARALLEL_NUM": 1,
+ "num_stages": 1,
+ "num_warps": 4
+ },
+ "128": {
+ "BLOCK_SEQ": 1,
+ "HEAD_PARALLEL_NUM": 8,
+ "num_stages": 5,
+ "num_warps": 2
+ },
+ "16": {
+ "BLOCK_SEQ": 1,
+ "HEAD_PARALLEL_NUM": 8,
+ "num_stages": 5,
+ "num_warps": 4
+ },
+ "2048": {
+ "BLOCK_SEQ": 1,
+ "HEAD_PARALLEL_NUM": 1,
+ "num_stages": 1,
+ "num_warps": 2
+ },
+ "256": {
+ "BLOCK_SEQ": 2,
+ "HEAD_PARALLEL_NUM": 8,
+ "num_stages": 5,
+ "num_warps": 2
+ },
+ "32": {
+ "BLOCK_SEQ": 1,
+ "HEAD_PARALLEL_NUM": 8,
+ "num_stages": 5,
+ "num_warps": 8
+ },
+ "4096": {
+ "BLOCK_SEQ": 1,
+ "HEAD_PARALLEL_NUM": 1,
+ "num_stages": 2,
+ "num_warps": 1
+ },
+ "64": {
+ "BLOCK_SEQ": 1,
+ "HEAD_PARALLEL_NUM": 8,
+ "num_stages": 1,
+ "num_warps": 4
+ },
+ "8": {
+ "BLOCK_SEQ": 1,
+ "HEAD_PARALLEL_NUM": 8,
+ "num_stages": 2,
+ "num_warps": 8
+ },
+ "8192": {
+ "BLOCK_SEQ": 2,
+ "HEAD_PARALLEL_NUM": 1,
+ "num_stages": 2,
+ "num_warps": 1
+ }
+}
\ No newline at end of file
diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H200/silu_and_mul_fwd:v1/{N=2048,out_dtype=torch.bfloat16}_NVIDIA_H200.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H200/silu_and_mul_fwd:v1/{N=2048,out_dtype=torch.bfloat16}_NVIDIA_H200.json
new file mode 100644
index 0000000000..d0ea86fac3
--- /dev/null
+++ b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H200/silu_and_mul_fwd:v1/{N=2048,out_dtype=torch.bfloat16}_NVIDIA_H200.json
@@ -0,0 +1,554 @@
+{
+ "1": {
+ "BLOCK_M": 256,
+ "BLOCK_N": 64,
+ "NUM_STAGES": 1,
+ "num_warps": 1
+ },
+ "100": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 4
+ },
+ "1024": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 1,
+ "num_warps": 4
+ },
+ "1152": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 1,
+ "num_warps": 4
+ },
+ "12160": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "128": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "1280": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 1,
+ "num_warps": 1
+ },
+ "13056": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "13312": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "13696": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "13824": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "13952": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "1408": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "14080": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "14336": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "14464": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "14720": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "14848": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "14976": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "15104": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "1536": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "16": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 8
+ },
+ "1664": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 1,
+ "num_warps": 1
+ },
+ "1920": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "2048": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "2176": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "2304": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "24192": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "2432": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "24960": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "25088": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "25344": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "256": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 1,
+ "num_warps": 1
+ },
+ "2560": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "25600": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "25856": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "25984": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "26112": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "26368": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "26624": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "2688": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "27136": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "27520": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "27776": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "2816": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "28160": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "28800": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "2944": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "3072": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "32": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 2,
+ "num_warps": 1
+ },
+ "3200": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "3328": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "3456": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "3584": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "3712": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "384": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "3840": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "3968": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "4096": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "4224": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "4352": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "4480": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "46336": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "46592": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "48896": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "49152": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "49280": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "50560": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "50688": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "50944": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "512": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "51328": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "52608": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "53248": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "53632": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "54400": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "54656": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "55040": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "64": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "640": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 1,
+ "num_warps": 4
+ },
+ "7296": {
+ "BLOCK_M": 32,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "7552": {
+ "BLOCK_M": 32,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "768": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "7808": {
+ "BLOCK_M": 32,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "7936": {
+ "BLOCK_M": 32,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "8": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 32,
+ "NUM_STAGES": 4,
+ "num_warps": 8
+ },
+ "8064": {
+ "BLOCK_M": 32,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "8192": {
+ "BLOCK_M": 32,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "8320": {
+ "BLOCK_M": 32,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "8448": {
+ "BLOCK_M": 32,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "8704": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "896": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ }
+}
\ No newline at end of file
diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H200/silu_and_mul_fwd:v1/{N=256,out_dtype=torch.bfloat16}_NVIDIA_H200.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H200/silu_and_mul_fwd:v1/{N=256,out_dtype=torch.bfloat16}_NVIDIA_H200.json
new file mode 100644
index 0000000000..fbd3649737
--- /dev/null
+++ b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H200/silu_and_mul_fwd:v1/{N=256,out_dtype=torch.bfloat16}_NVIDIA_H200.json
@@ -0,0 +1,146 @@
+{
+ "1": {
+ "BLOCK_M": 128,
+ "BLOCK_N": 32,
+ "NUM_STAGES": 2,
+ "num_warps": 4
+ },
+ "100": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 4
+ },
+ "1024": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 1,
+ "num_warps": 1
+ },
+ "12288": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 1,
+ "num_warps": 1
+ },
+ "128": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 1,
+ "num_warps": 1
+ },
+ "1536": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 1,
+ "num_warps": 4
+ },
+ "16": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 32,
+ "NUM_STAGES": 2,
+ "num_warps": 4
+ },
+ "192": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 8
+ },
+ "2048": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 1,
+ "num_warps": 4
+ },
+ "24576": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "256": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 1,
+ "num_warps": 1
+ },
+ "32": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 8
+ },
+ "384": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 2,
+ "num_warps": 1
+ },
+ "4096": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 4
+ },
+ "48": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 32,
+ "NUM_STAGES": 2,
+ "num_warps": 4
+ },
+ "49152": {
+ "BLOCK_M": 32,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 4,
+ "num_warps": 1
+ },
+ "6": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 32,
+ "NUM_STAGES": 4,
+ "num_warps": 8
+ },
+ "600": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 1,
+ "num_warps": 4
+ },
+ "6144": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 1,
+ "num_warps": 4
+ },
+ "64": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 4
+ },
+ "768": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 1,
+ "num_warps": 4
+ },
+ "8": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 64,
+ "NUM_STAGES": 1,
+ "num_warps": 1
+ },
+ "8192": {
+ "BLOCK_M": 8,
+ "BLOCK_N": 256,
+ "NUM_STAGES": 1,
+ "num_warps": 4
+ },
+ "96": {
+ "BLOCK_M": 1,
+ "BLOCK_N": 128,
+ "NUM_STAGES": 4,
+ "num_warps": 8
+ }
+}
\ No newline at end of file
diff --git a/lightllm/models/__init__.py b/lightllm/models/__init__.py
index f619b1d88f..3d376d160d 100644
--- a/lightllm/models/__init__.py
+++ b/lightllm/models/__init__.py
@@ -20,6 +20,7 @@
from lightllm.models.phi3.model import Phi3TpPartModel
from lightllm.models.deepseek2.model import Deepseek2TpPartModel
from lightllm.models.deepseek3_2.model import Deepseek3_2TpPartModel
+from lightllm.models.deepseek_v4.model import DeepseekV4TpPartModel
from lightllm.models.glm4_moe_lite.model import Glm4MoeLiteTpPartModel
from lightllm.models.internvl.model import (
InternVLLlamaTpPartModel,
diff --git a/lightllm/models/deepseek2/triton_kernel/rotary_emb.py b/lightllm/models/deepseek2/triton_kernel/rotary_emb.py
index 30e5a59248..a8f851de2a 100644
--- a/lightllm/models/deepseek2/triton_kernel/rotary_emb.py
+++ b/lightllm/models/deepseek2/triton_kernel/rotary_emb.py
@@ -29,6 +29,8 @@ def _rotary_kernel(
BLOCK_SEQ: tl.constexpr,
BLOCK_DMODEL: tl.constexpr,
NUM_STAGE: tl.constexpr,
+ HAS_K: tl.constexpr,
+ INVERSE: tl.constexpr,
):
head_start_index = tl.program_id(0)
seq_block_index = tl.program_id(1)
@@ -44,6 +46,8 @@ def _rotary_kernel(
off_dimcos_sin = seq_index * stride_cosbs + cos_range * stride_cosd
cos = tl.load(Cos + off_dimcos_sin)
sin = tl.load(Sin + off_dimcos_sin)
+ if INVERSE:
+ sin = -sin
if HEAD_PARALLEL_NUM == 1:
for q_head_index in tl.static_range(0, HEAD_Q, step=1):
@@ -56,18 +60,19 @@ def _rotary_kernel(
tl.store(Q + off_q0, out_q0)
tl.store(Q + off_q1, out_q1)
- for k_head_index in tl.static_range(0, HEAD_K, step=1):
- off_k0 = seq_index * stride_kbs + k_head_index * stride_kh + dim_range0 * stride_kd
- off_k1 = seq_index * stride_kbs + k_head_index * stride_kh + dim_range1 * stride_kd
+ if HAS_K:
+ for k_head_index in tl.static_range(0, HEAD_K, step=1):
+ off_k0 = seq_index * stride_kbs + k_head_index * stride_kh + dim_range0 * stride_kd
+ off_k1 = seq_index * stride_kbs + k_head_index * stride_kh + dim_range1 * stride_kd
- k0 = tl.load(K + off_k0)
- k1 = tl.load(K + off_k1)
+ k0 = tl.load(K + off_k0)
+ k1 = tl.load(K + off_k1)
- out_k0 = k0 * cos - k1 * sin
- out_k1 = k0 * sin + k1 * cos
+ out_k0 = k0 * cos - k1 * sin
+ out_k1 = k0 * sin + k1 * cos
- tl.store(K + off_k0, out_k0)
- tl.store(K + off_k1, out_k1)
+ tl.store(K + off_k0, out_k0)
+ tl.store(K + off_k1, out_k1)
else:
for q_head_index in tl.range(head_start_index, HEAD_Q, step=HEAD_PARALLEL_NUM, num_stages=NUM_STAGE):
off_q0 = seq_index * stride_qbs + q_head_index * stride_qh + dim_range0 * stride_qd
@@ -79,18 +84,19 @@ def _rotary_kernel(
tl.store(Q + off_q0, out_q0)
tl.store(Q + off_q1, out_q1)
- for k_head_index in tl.range(head_start_index, HEAD_K, step=HEAD_PARALLEL_NUM, num_stages=NUM_STAGE):
- off_k0 = seq_index * stride_kbs + k_head_index * stride_kh + dim_range0 * stride_kd
- off_k1 = seq_index * stride_kbs + k_head_index * stride_kh + dim_range1 * stride_kd
+ if HAS_K:
+ for k_head_index in tl.range(head_start_index, HEAD_K, step=HEAD_PARALLEL_NUM, num_stages=NUM_STAGE):
+ off_k0 = seq_index * stride_kbs + k_head_index * stride_kh + dim_range0 * stride_kd
+ off_k1 = seq_index * stride_kbs + k_head_index * stride_kh + dim_range1 * stride_kd
- k0 = tl.load(K + off_k0)
- k1 = tl.load(K + off_k1)
+ k0 = tl.load(K + off_k0)
+ k1 = tl.load(K + off_k1)
- out_k0 = k0 * cos - k1 * sin
- out_k1 = k0 * sin + k1 * cos
+ out_k0 = k0 * cos - k1 * sin
+ out_k1 = k0 * sin + k1 * cos
- tl.store(K + off_k0, out_k0)
- tl.store(K + off_k1, out_k1)
+ tl.store(K + off_k0, out_k0)
+ tl.store(K + off_k1, out_k1)
return
@@ -109,7 +115,10 @@ def get_test_configs():
def get_static_key(q, k):
- head_num_q, head_num_k, head_dim = q.shape[1], k.shape[1], q.shape[2]
+ assert q is not None, "q can not be None"
+ head_num_q = q.shape[1]
+ head_num_k = k.shape[1] if k is not None else 0
+ head_dim = q.shape[2]
return {
"Q_HEAD_NUM": head_num_q,
"K_HEAD_NUM": head_num_k,
@@ -126,12 +135,17 @@ def get_static_key(q, k):
mutates_args=["q", "k"],
)
@torch.no_grad()
-def rotary_emb_fwd(q, k, cos, sin, run_config=None):
+def rotary_emb_fwd(q, k, cos, sin, inverse=False, run_config=None):
+ assert q is not None, "q can not be None"
+ has_k = k is not None and k.shape[1] != 0
total_len = q.shape[0]
- head_num_q, head_num_k = q.shape[1], k.shape[1]
+ head_num_q = q.shape[1]
+ head_num_k = k.shape[1] if k is not None else 0
head_dim = q.shape[2]
assert q.shape[0] == cos.shape[0] and q.shape[0] == sin.shape[0], f"q shape {q.shape} cos shape {cos.shape}"
- assert k.shape[0] == cos.shape[0] and k.shape[0] == sin.shape[0], f"k shape {k.shape} cos shape {cos.shape}"
+ if k is not None:
+ assert k.shape[0] == cos.shape[0] and k.shape[0] == sin.shape[0], f"k shape {k.shape} cos shape {cos.shape}"
+ assert k.shape[2] == head_dim, f"k shape {k.shape} q head_dim {head_dim}"
assert triton.next_power_of_2(head_dim) == head_dim
if not run_config:
@@ -157,9 +171,9 @@ def rotary_emb_fwd(q, k, cos, sin, run_config=None):
stride_qbs=q.stride(0),
stride_qh=q.stride(1),
stride_qd=q.stride(2),
- stride_kbs=k.stride(0),
- stride_kh=k.stride(1),
- stride_kd=k.stride(2),
+ stride_kbs=k.stride(0) if k is not None else 0,
+ stride_kh=k.stride(1) if k is not None else 0,
+ stride_kd=k.stride(2) if k is not None else 0,
stride_cosbs=cos.stride(0),
stride_cosd=cos.stride(1),
stride_sinbs=sin.stride(0),
@@ -171,6 +185,8 @@ def rotary_emb_fwd(q, k, cos, sin, run_config=None):
BLOCK_SEQ=BLOCK_SEQ,
BLOCK_DMODEL=head_dim,
NUM_STAGE=num_stages,
+ HAS_K=has_k,
+ INVERSE=inverse,
num_warps=num_warps,
num_stages=num_stages,
)
diff --git a/lightllm/models/deepseek_v4/__init__.py b/lightllm/models/deepseek_v4/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/lightllm/models/deepseek_v4/encoding/__init__.py b/lightllm/models/deepseek_v4/encoding/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/lightllm/models/deepseek_v4/encoding/encoding_dsv4.py b/lightllm/models/deepseek_v4/encoding/encoding_dsv4.py
new file mode 100644
index 0000000000..6cbd5f9bfb
--- /dev/null
+++ b/lightllm/models/deepseek_v4/encoding/encoding_dsv4.py
@@ -0,0 +1,762 @@
+"""
+DeepSeek-V4 Encoding
+
+A self-contained implementation for encoding/decoding DeepSeek-V4 chat messages
+with tool calling, thinking mode, and quick instruction task support.
+"""
+
+from typing import Any, Dict, List, Union, Optional, Tuple
+import copy
+import json
+import re
+
+# ============================================================
+# Special Tokens
+# ============================================================
+
+bos_token: str = "<｜begin▁of▁sentence｜>"
+eos_token: str = "<｜end▁of▁sentence｜>"
+thinking_start_token: str = ""
+thinking_end_token: str = ""
+dsml_token: str = "｜DSML｜"
+
+USER_SP_TOKEN = "<｜User｜>"
+ASSISTANT_SP_TOKEN = "<｜Assistant｜>"
+LATEST_REMINDER_SP_TOKEN = "<｜latest_reminder｜>"
+
+# Task special tokens for internal classification tasks
+DS_TASK_SP_TOKENS = {
+ "action": "<｜action｜>",
+ "query": "<｜query｜>",
+ "authority": "<｜authority｜>",
+ "domain": "<｜domain｜>",
+ "title": "<｜title｜>",
+ "read_url": "<｜read_url｜>",
+}
+VALID_TASKS = set(DS_TASK_SP_TOKENS.keys())
+
+# ============================================================
+# Templates
+# ============================================================
+
+system_msg_template: str = "{content}"
+user_msg_template: str = "{content}"
+latest_reminder_msg_template: str = "{content}"
+assistant_msg_template: str = "{reasoning}{content}{tool_calls}" + eos_token
+assistant_msg_wo_eos_template: str = "{reasoning}{content}{tool_calls}"
+thinking_template: str = "{reasoning_content}"
+
+response_format_template: str = (
+ "## Response Format:\n\nYou MUST strictly adhere to the following schema to reply:\n{schema}"
+)
+tool_call_template: str = '<{dsml_token}invoke name="{name}">\n{arguments}\n'
+tool_calls_template = "<{dsml_token}{tc_block_name}>\n{tool_calls}\n"
+tool_calls_block_name: str = "tool_calls"
+
+tool_output_template: str = "{content}"
+
+REASONING_EFFORT_MAX = (
+ "Reasoning Effort: Absolute maximum with no shortcuts permitted.\n"
+ "You MUST be very thorough in your thinking and comprehensively decompose the problem to resolve the root cause, rigorously stress-testing your logic against all potential paths, edge cases, and adversarial scenarios.\n" # noqa: E501
+ "Explicitly write out your entire deliberation process, documenting every intermediate step, considered alternative, and rejected hypothesis to ensure absolutely no assumption is left unchecked.\n\n" # noqa: E501
+)
+
+TOOLS_TEMPLATE = """## Tools
+
+You have access to a set of tools to help answer the user's question. You can invoke tools by writing a \
+"<{dsml_token}tool_calls>" block like the following:
+
+<{dsml_token}tool_calls>
+<{dsml_token}invoke name="$TOOL_NAME">
+<{dsml_token}parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE
+...
+
+<{dsml_token}invoke name="$TOOL_NAME2">
+...
+
+
+
+String parameters should be specified as is and set `string="true"`. For all other types (numbers, \
+booleans, arrays, objects), pass the value in JSON format and set `string="false"`.
+
+If thinking_mode is enabled (triggered by {thinking_start_token}), you MUST output your complete \
+reasoning inside {thinking_start_token}...{thinking_end_token} BEFORE any tool calls or final response.
+
+Otherwise, output directly after {thinking_end_token} with tool calls or final response.
+
+### Available Tool Schemas
+
+{tool_schemas}
+
+You MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls.
+"""
+
+# ============================================================
+# Utility Functions
+# ============================================================
+
+
+def to_json(value: Any) -> str:
+ """Serialize a value to JSON string."""
+ try:
+ return json.dumps(value, ensure_ascii=False)
+ except:
+ return json.dumps(value, ensure_ascii=True)
+
+
+def tools_from_openai_format(tools):
+ """Extract function definitions from OpenAI-format tool list."""
+ return [tool["function"] for tool in tools]
+
+
+def tool_calls_from_openai_format(tool_calls):
+ """Convert OpenAI-format tool calls to internal format."""
+ return [
+ {
+ "name": tool_call["function"]["name"],
+ "arguments": tool_call["function"]["arguments"],
+ }
+ for tool_call in tool_calls
+ ]
+
+
+def tool_calls_to_openai_format(tool_calls):
+ """Convert internal tool calls to OpenAI format."""
+ return [
+ {
+ "type": "function",
+ "function": {
+ "name": tool_call["name"],
+ "arguments": tool_call["arguments"],
+ },
+ }
+ for tool_call in tool_calls
+ ]
+
+
+def encode_arguments_to_dsml(tool_call: Dict[str, str]) -> str:
+ """
+ Encode tool call arguments into DSML parameter format.
+
+ Args:
+ tool_call: Dict with "name" and "arguments" (JSON string) keys.
+
+ Returns:
+ DSML-formatted parameter string.
+ """
+ p_dsml_template = '<{dsml_token}parameter name="{key}" string="{is_str}">{value}'
+ P_dsml_strs = []
+
+ try:
+ arguments = json.loads(tool_call["arguments"])
+ except Exception:
+ arguments = {"arguments": tool_call["arguments"]}
+
+ for k, v in arguments.items():
+ p_dsml_str = p_dsml_template.format(
+ dsml_token=dsml_token,
+ key=k,
+ is_str="true" if isinstance(v, str) else "false",
+ value=v if isinstance(v, str) else to_json(v),
+ )
+ P_dsml_strs.append(p_dsml_str)
+
+ return "\n".join(P_dsml_strs)
+
+
+def decode_dsml_to_arguments(tool_name: str, tool_args: Dict[str, Tuple[str, str]]) -> Dict[str, str]:
+ """
+ Decode DSML parameters back to a tool call dict.
+
+ Args:
+ tool_name: Name of the tool.
+ tool_args: Dict mapping param_name -> (value, is_string_flag).
+
+ Returns:
+ Dict with "name" and "arguments" (JSON string) keys.
+ """
+
+ def _decode_value(key: str, value: str, string: str):
+ if string == "true":
+ value = to_json(value)
+ return f"{to_json(key)}: {value}"
+
+ tool_args_json = "{" + ", ".join([_decode_value(k, v, string=is_str) for k, (v, is_str) in tool_args.items()]) + "}"
+ return dict(name=tool_name, arguments=tool_args_json)
+
+
+def render_tools(tools: List[Dict[str, Union[str, Dict[str, Any]]]]) -> str:
+ """
+ Render tool schemas into the system prompt format.
+
+ Args:
+ tools: List of tool schema dicts (each with name, description, parameters).
+
+ Returns:
+ Formatted tools section string.
+ """
+ tools_json = [to_json(t) for t in tools]
+
+ return TOOLS_TEMPLATE.format(
+ tool_schemas="\n".join(tools_json),
+ dsml_token=dsml_token,
+ thinking_start_token=thinking_start_token,
+ thinking_end_token=thinking_end_token,
+ )
+
+
+def find_last_user_index(messages: List[Dict[str, Any]]) -> int:
+ """Find the index of the last user/developer message."""
+ last_user_index = -1
+ for idx in range(len(messages) - 1, -1, -1):
+ if messages[idx].get("role") in ["user", "developer"]:
+ last_user_index = idx
+ break
+ return last_user_index
+
+
+# ============================================================
+# Message Rendering
+# ============================================================
+
+
+def render_message(
+ index: int,
+ messages: List[Dict[str, Any]],
+ thinking_mode: str,
+ drop_thinking: bool = True,
+ reasoning_effort: Optional[str] = None,
+) -> str:
+ """
+ Render a single message at the given index into its encoded string form.
+
+ This is the core function that converts each message in the conversation
+ into the DeepSeek-V4 format.
+
+ Args:
+ index: Index of the message to render.
+ messages: Full list of messages in the conversation.
+ thinking_mode: Either "chat" or "thinking".
+ drop_thinking: Whether to drop reasoning content from earlier turns.
+ reasoning_effort: Optional reasoning effort level ("max", "high", or None).
+
+ Returns:
+ Encoded string for this message.
+ """
+ assert 0 <= index < len(messages)
+ assert thinking_mode in ["chat", "thinking"], f"Invalid thinking_mode `{thinking_mode}`"
+
+ prompt = ""
+ msg = messages[index]
+ last_user_idx = find_last_user_index(messages)
+
+ role = msg.get("role")
+ content = msg.get("content")
+ tools = msg.get("tools")
+ response_format = msg.get("response_format")
+ tool_calls = msg.get("tool_calls")
+ reasoning_content = msg.get("reasoning_content")
+ wo_eos = msg.get("wo_eos", False)
+
+ if tools:
+ tools = tools_from_openai_format(tools)
+ if tool_calls:
+ tool_calls = tool_calls_from_openai_format(tool_calls)
+
+ # Reasoning effort prefix (only at index 0 in thinking mode with max effort)
+ assert reasoning_effort in ["max", None, "high"], f"Invalid reasoning effort: {reasoning_effort}"
+ if index == 0 and thinking_mode == "thinking" and reasoning_effort == "max":
+ prompt += REASONING_EFFORT_MAX
+
+ if role == "system":
+ prompt += system_msg_template.format(content=content or "")
+ if tools:
+ prompt += "\n\n" + render_tools(tools)
+ if response_format:
+ prompt += "\n\n" + response_format_template.format(schema=to_json(response_format))
+
+ elif role == "developer":
+ assert content, f"Invalid message for role `{role}`: {msg}"
+
+ content_developer = USER_SP_TOKEN
+ content_developer += content
+
+ if tools:
+ content_developer += "\n\n" + render_tools(tools)
+ if response_format:
+ content_developer += "\n\n" + response_format_template.format(schema=to_json(response_format))
+
+ prompt += user_msg_template.format(content=content_developer)
+
+ elif role == "user":
+ prompt += USER_SP_TOKEN
+
+ # Handle content blocks (tool results mixed with text)
+ content_blocks = msg.get("content_blocks")
+ if content_blocks:
+ parts = []
+ for block in content_blocks:
+ block_type = block.get("type")
+ if block_type == "text":
+ parts.append(block.get("text", ""))
+ elif block_type == "tool_result":
+ tool_content = block.get("content", "")
+ if isinstance(tool_content, list):
+ text_parts = []
+ for b in tool_content:
+ if b.get("type") == "text":
+ text_parts.append(b.get("text", ""))
+ else:
+ text_parts.append(f"[Unsupported {b.get('type')}]")
+ tool_content = "\n\n".join(text_parts)
+ parts.append(tool_output_template.format(content=tool_content))
+ else:
+ parts.append(f"[Unsupported {block_type}]")
+ prompt += "\n\n".join(parts)
+ else:
+ prompt += content or ""
+
+ elif role == "latest_reminder":
+ prompt += LATEST_REMINDER_SP_TOKEN + latest_reminder_msg_template.format(content=content)
+
+ elif role == "tool":
+ raise NotImplementedError(
+ "deepseek_v4 merges tool messages into user; please preprocess with merge_tool_messages()"
+ )
+
+ elif role == "assistant":
+ thinking_part = ""
+ tc_content = ""
+
+ if tool_calls:
+ tc_list = [
+ tool_call_template.format(
+ dsml_token=dsml_token, name=tc.get("name"), arguments=encode_arguments_to_dsml(tc)
+ )
+ for tc in tool_calls
+ ]
+ tc_content += "\n\n" + tool_calls_template.format(
+ dsml_token=dsml_token,
+ tool_calls="\n".join(tc_list),
+ tc_block_name=tool_calls_block_name,
+ )
+
+ summary_content = content or ""
+ rc = reasoning_content or ""
+
+ # Check if previous message has a task - if so, this is a task output (no thinking)
+ prev_has_task = index - 1 >= 0 and messages[index - 1].get("task") is not None
+
+ if thinking_mode == "thinking" and not prev_has_task:
+ if not drop_thinking or index > last_user_idx:
+ thinking_part = thinking_template.format(reasoning_content=rc) + thinking_end_token
+ else:
+ thinking_part = ""
+
+ if wo_eos:
+ prompt += assistant_msg_wo_eos_template.format(
+ reasoning=thinking_part,
+ content=summary_content,
+ tool_calls=tc_content,
+ )
+ else:
+ prompt += assistant_msg_template.format(
+ reasoning=thinking_part,
+ content=summary_content,
+ tool_calls=tc_content,
+ )
+ else:
+ raise NotImplementedError(f"Unknown role: {role}")
+
+ # Append transition tokens based on what follows
+ if index + 1 < len(messages) and messages[index + 1].get("role") not in ["assistant", "latest_reminder"]:
+ return prompt
+
+ task = messages[index].get("task")
+ if task is not None:
+ # Task special token for internal classification tasks
+ assert task in VALID_TASKS, f"Invalid task: '{task}'. Valid tasks are: {list(VALID_TASKS)}"
+ task_sp_token = DS_TASK_SP_TOKENS[task]
+
+ if task != "action":
+ # Non-action tasks: append task sp token directly after the message
+ prompt += task_sp_token
+ else:
+ # Action task: append Assistant + thinking token + action sp token
+ prompt += ASSISTANT_SP_TOKEN
+ prompt += thinking_end_token if thinking_mode != "thinking" else thinking_start_token
+ prompt += task_sp_token
+
+ elif messages[index].get("role") in ["user", "developer"]:
+ # Normal generation: append Assistant + thinking token
+ prompt += ASSISTANT_SP_TOKEN
+ if not drop_thinking and thinking_mode == "thinking":
+ prompt += thinking_start_token
+ elif drop_thinking and thinking_mode == "thinking" and index >= last_user_idx:
+ prompt += thinking_start_token
+ else:
+ prompt += thinking_end_token
+
+ return prompt
+
+
+# ============================================================
+# Preprocessing
+# ============================================================
+
+
+def merge_tool_messages(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+ """
+ Merge tool messages into the preceding user message using content_blocks format.
+
+ DeepSeek-V4 does not have a standalone "tool" role; instead, tool results
+ are encoded as blocks within user messages.
+
+ This function converts a standard OpenAI-format conversation (with separate
+ "tool" role messages) into V4 format where tool results are merged into
+ user messages.
+
+ Args:
+ messages: List of message dicts in OpenAI format.
+
+ Returns:
+ Processed message list with tool messages merged into user messages.
+ """
+ merged: List[Dict[str, Any]] = []
+
+ for msg in messages:
+ msg = copy.deepcopy(msg)
+ role = msg.get("role")
+
+ if role == "tool":
+ # Convert tool message to a user message with tool_result block
+ tool_block = {
+ "type": "tool_result",
+ "tool_use_id": msg.get("tool_call_id", ""),
+ "content": msg.get("content", ""),
+ }
+ # Merge into previous message if it's already a user (merged tool)
+ if merged and merged[-1].get("role") == "user" and "content_blocks" in merged[-1]:
+ merged[-1]["content_blocks"].append(tool_block)
+ else:
+ merged.append(
+ {
+ "role": "user",
+ "content_blocks": [tool_block],
+ }
+ )
+ elif role == "user":
+ text_block = {"type": "text", "text": msg.get("content", "")}
+ if (
+ merged
+ and merged[-1].get("role") == "user"
+ and "content_blocks" in merged[-1]
+ and merged[-1].get("task") is None
+ ):
+ merged[-1]["content_blocks"].append(text_block)
+ else:
+ new_msg = {
+ "role": "user",
+ "content": msg.get("content", ""),
+ "content_blocks": [text_block],
+ }
+ # Preserve extra fields (task, wo_eos, mask, etc.)
+ for key in ("task", "wo_eos", "mask"):
+ if key in msg:
+ new_msg[key] = msg[key]
+ merged.append(new_msg)
+ else:
+ merged.append(msg)
+
+ return merged
+
+
+def sort_tool_results_by_call_order(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+ """
+ Sort tool_result blocks within user messages by the order of tool_calls
+ in the preceding assistant message.
+
+ Args:
+ messages: Preprocessed message list (after merge_tool_messages).
+
+ Returns:
+ Message list with sorted tool result blocks.
+ """
+ last_tool_call_order: Dict[str, int] = {}
+
+ for msg in messages:
+ role = msg.get("role")
+ if role == "assistant" and msg.get("tool_calls"):
+ last_tool_call_order = {}
+ for idx, tc in enumerate(msg["tool_calls"]):
+ tc_id = tc.get("id") or tc.get("function", {}).get("id", "")
+ if tc_id:
+ last_tool_call_order[tc_id] = idx
+
+ elif role == "user" and msg.get("content_blocks"):
+ tool_blocks = [b for b in msg["content_blocks"] if b.get("type") == "tool_result"]
+ if len(tool_blocks) > 1 and last_tool_call_order:
+ sorted_blocks = sorted(tool_blocks, key=lambda b: last_tool_call_order.get(b.get("tool_use_id", ""), 0))
+ sorted_idx = 0
+ new_blocks = []
+ for block in msg["content_blocks"]:
+ if block.get("type") == "tool_result":
+ new_blocks.append(sorted_blocks[sorted_idx])
+ sorted_idx += 1
+ else:
+ new_blocks.append(block)
+ msg["content_blocks"] = new_blocks
+
+ return messages
+
+
+# ============================================================
+# Main Encoding Function
+# ============================================================
+
+
+def encode_messages(
+ messages: List[Dict[str, Any]],
+ thinking_mode: str,
+ context: Optional[List[Dict[str, Any]]] = None,
+ drop_thinking: bool = True,
+ add_default_bos_token: bool = True,
+ reasoning_effort: Optional[str] = None,
+) -> str:
+ """
+ Encode a list of messages into the DeepSeek-V4 prompt format.
+
+ This is the main entry point for encoding conversations. It handles:
+ - BOS token insertion
+ - Thinking mode with optional reasoning content dropping
+ - Tool message merging into user messages
+ - Multi-turn conversation context
+
+ Args:
+ messages: List of message dicts to encode.
+ thinking_mode: Either "chat" or "thinking".
+ context: Optional preceding context messages (already encoded prefix).
+ drop_thinking: If True, drop reasoning_content from earlier assistant turns
+ (only keep reasoning for messages after the last user message).
+ add_default_bos_token: Whether to prepend BOS token at conversation start.
+ reasoning_effort: Optional reasoning effort level ("max", "high", or None).
+
+ Returns:
+ The encoded prompt string.
+ """
+ context = context if context else []
+
+ # Preprocess: merge tool messages and sort tool results
+ messages = merge_tool_messages(messages)
+ messages = sort_tool_results_by_call_order(context + messages)[len(context) :]
+ if context:
+ context = merge_tool_messages(context)
+ context = sort_tool_results_by_call_order(context)
+
+ full_messages = context + messages
+
+ prompt = bos_token if add_default_bos_token and len(context) == 0 else ""
+
+ # Resolve drop_thinking: if any message has tools defined, don't drop thinking
+ effective_drop_thinking = drop_thinking
+ if any(m.get("tools") for m in full_messages):
+ effective_drop_thinking = False
+
+ if thinking_mode == "thinking" and effective_drop_thinking:
+ full_messages = _drop_thinking_messages(full_messages)
+ # After dropping, recalculate how many messages to render
+ # (context may have shrunk too)
+ num_to_render = len(full_messages) - len(_drop_thinking_messages(context))
+ context_len = len(full_messages) - num_to_render
+ else:
+ num_to_render = len(messages)
+ context_len = len(context)
+
+ for idx in range(num_to_render):
+ prompt += render_message(
+ idx + context_len,
+ full_messages,
+ thinking_mode=thinking_mode,
+ drop_thinking=effective_drop_thinking,
+ reasoning_effort=reasoning_effort,
+ )
+
+ return prompt
+
+
+def _drop_thinking_messages(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+ """
+ Drop reasoning_content and non-essential messages before the last user message.
+
+ Behavior:
+ - Messages with role in ["user", "system", "tool", "latest_reminder"] are always kept.
+ - Messages at or after the last user index are always kept.
+ - Assistant messages before the last user get reasoning_content removed.
+ - Developer messages before the last user are dropped entirely.
+ """
+ last_user_idx = find_last_user_index(messages)
+ result = []
+ keep_roles = {"user", "system", "tool", "latest_reminder", "direct_search_results"}
+
+ for idx, msg in enumerate(messages):
+ role = msg.get("role")
+ if role in keep_roles or idx >= last_user_idx:
+ result.append(msg)
+ elif role == "assistant":
+ msg = copy.copy(msg)
+ msg.pop("reasoning_content", None)
+ result.append(msg)
+ # developer and other roles before last_user_idx are dropped
+
+ return result
+
+
+# ============================================================
+# Parsing (Decoding model output)
+# ============================================================
+
+
+def _read_until_stop(index: int, text: str, stop: List[str]) -> Tuple[int, str, Optional[str]]:
+ """
+ Read text from index until one of the stop strings is found.
+
+ Returns:
+ Tuple of (new_index, content_before_stop, matched_stop_string_or_None).
+ """
+ min_pos = len(text)
+ matched_stop = None
+
+ for s in stop:
+ pos = text.find(s, index)
+ if pos != -1 and pos < min_pos:
+ min_pos = pos
+ matched_stop = s
+
+ if matched_stop:
+ content = text[index:min_pos]
+ return min_pos + len(matched_stop), content, matched_stop
+ else:
+ content = text[index:]
+ return len(text), content, None
+
+
+def parse_tool_calls(index: int, text: str) -> Tuple[int, Optional[str], List[Dict[str, str]]]:
+ """
+ Parse DSML tool calls from text starting at the given index.
+
+ Args:
+ index: Starting position in text.
+ text: The full text to parse.
+
+ Returns:
+ Tuple of (new_index, last_stop_token, list_of_tool_call_dicts).
+ Each tool call dict has "name" and "arguments" keys.
+ """
+ tool_calls: List[Dict[str, Any]] = []
+ stop_token = None
+ tool_calls_end_token = f""
+
+ while index < len(text):
+ index, _, stop_token = _read_until_stop(index, text, [f"<{dsml_token}invoke", tool_calls_end_token])
+ if _ != ">\n":
+ raise ValueError(f"Tool call format error: expected '>\\n' but got '{_}'")
+
+ if stop_token == tool_calls_end_token:
+ break
+
+ if stop_token is None:
+ raise ValueError("Missing special token in tool calls")
+
+ index, tool_name_content, stop_token = _read_until_stop(
+ index, text, [f"<{dsml_token}parameter", f"\n$', tool_name_content, flags=re.DOTALL)
+ if len(p_tool_name) != 1:
+ raise ValueError(f"Tool name format error: '{tool_name_content}'")
+ tool_name = p_tool_name[0]
+
+ tool_args: Dict[str, Tuple[str, str]] = {}
+ while stop_token == f"<{dsml_token}parameter":
+ index, param_content, stop_token = _read_until_stop(index, text, [f"/{dsml_token}parameter"])
+
+ param_kv = re.findall(r'^ name="(.*?)" string="(true|false)">(.*?)<$', param_content, flags=re.DOTALL)
+ if len(param_kv) != 1:
+ raise ValueError(f"Parameter format error: '{param_content}'")
+ param_name, string, param_value = param_kv[0]
+
+ if param_name in tool_args:
+ raise ValueError(f"Duplicate parameter name: '{param_name}'")
+ tool_args[param_name] = (param_value, string)
+
+ index, content, stop_token = _read_until_stop(
+ index, text, [f"<{dsml_token}parameter", f"\n":
+ raise ValueError(f"Parameter format error: expected '>\\n' but got '{content}'")
+
+ tool_call = decode_dsml_to_arguments(tool_name=tool_name, tool_args=tool_args)
+ tool_calls.append(tool_call)
+
+ return index, stop_token, tool_calls
+
+
+def parse_message_from_completion_text(text: str, thinking_mode: str) -> Dict[str, Any]:
+ """
+ Parse a model completion text into a structured assistant message.
+
+ This function takes the raw text output from the model (a single assistant turn)
+ and extracts:
+ - reasoning_content (thinking block)
+ - content (summary/response)
+ - tool_calls (if any)
+
+ NOTE: This function is designed to parse only correctly formatted strings and
+ will raise ValueError for malformed output.
+
+ Args:
+ text: The raw completion text (including EOS token).
+ thinking_mode: Either "chat" or "thinking".
+
+ Returns:
+ Dict with keys: "role", "content", "reasoning_content", "tool_calls".
+ tool_calls are in OpenAI format.
+ """
+ summary_content, reasoning_content, tool_calls = "", "", []
+ index, stop_token = 0, None
+ tool_calls_start_token = f"\n\n<{dsml_token}{tool_calls_block_name}"
+
+ is_thinking = thinking_mode == "thinking"
+ is_tool_calling = False
+
+ if is_thinking:
+ index, content_delta, stop_token = _read_until_stop(index, text, [thinking_end_token, tool_calls_start_token])
+ reasoning_content = content_delta
+ assert stop_token == thinking_end_token, "Invalid thinking format: missing "
+
+ index, content_delta, stop_token = _read_until_stop(index, text, [eos_token, tool_calls_start_token])
+ summary_content = content_delta
+ if stop_token == tool_calls_start_token:
+ is_tool_calling = True
+ else:
+ assert stop_token == eos_token, "Invalid format: missing EOS token"
+
+ if is_tool_calling:
+ index, stop_token, tool_calls = parse_tool_calls(index, text)
+
+ index, tool_ends_text, stop_token = _read_until_stop(index, text, [eos_token])
+ assert not tool_ends_text, "Unexpected content after tool calls"
+
+ assert len(text) == index and stop_token in [eos_token, None], "Unexpected content at end"
+
+ for sp_token in [bos_token, eos_token, thinking_start_token, thinking_end_token, dsml_token]:
+ assert (
+ sp_token not in summary_content and sp_token not in reasoning_content
+ ), f"Unexpected special token '{sp_token}' in content"
+
+ return {
+ "role": "assistant",
+ "content": summary_content,
+ "reasoning_content": reasoning_content,
+ "tool_calls": tool_calls_to_openai_format(tool_calls),
+ }
diff --git a/lightllm/models/deepseek_v4/infer_struct.py b/lightllm/models/deepseek_v4/infer_struct.py
new file mode 100644
index 0000000000..e14e84c968
--- /dev/null
+++ b/lightllm/models/deepseek_v4/infer_struct.py
@@ -0,0 +1,98 @@
+import torch
+from lightllm.common.basemodel import InferStateInfo
+from lightllm.common.req_manager import DeepseekV4ReqManager
+from lightllm.common.kv_cache_mem_manager import DeepseekV4MemoryManager
+
+
+class DeepseekV4InferStateInfo(InferStateInfo):
+ req_manager: DeepseekV4ReqManager
+ mem_manager: DeepseekV4MemoryManager
+
+ """Per-token interleaved-rope cos/sin for the two rope variants (sliding / compressed), following
+ the gemma4 two-variant convention (_cos_cached_* -> position_cos_*). The full rope tables are
+ model constants and live on the model / layer infers, not here."""
+
+ def __init__(self):
+ super().__init__()
+ self.position_cos_sliding = None
+ self.position_sin_sliding = None
+ self.position_cos_compress = None
+ self.position_sin_compress = None
+ # layer-independent sparse-index metadata, built once per forward in init_some_extra_state
+ # (None until then so copy_for_cuda_graph's tensor-attr loop skips them).
+ self.dsv4_sparse_req_idx = None
+ self.dsv4_swa_indices = None
+ self.dsv4_swa_lengths = None
+ self.dsv4_c128_indices = None
+ self.dsv4_c128_lengths = None
+ self.dsv4_workspace = None
+ # token -> batch-position map for the compressor; built per prefill forward in init_some_extra_state.
+ self._dsv4_token_to_batch_idx = None
+ # lazily-built (first c4 layer) cache of layer-independent paged-c4 metadata; reused by the
+ # other c4 layers in the same forward. Plain tuple (not a tensor attr) so copy_for_cuda_graph
+ # ignores it -- it's a capture-time wiring of layer0->others, not a staged graph input.
+ self._c4_paged_meta = None
+
+ def _dsv4_index_max_kv_seq_len(self, model):
+ if (
+ not self.is_prefill
+ and model.graph is not None
+ and model.graph.can_run(self.batch_size, self.max_kv_seq_len)
+ ):
+ return model.graph.graph_max_len_in_batch
+ return self.max_kv_seq_len
+
+ def init_some_extra_state(self, model):
+ self._c4_paged_meta = None # reset per forward before any c4 layer runs
+ super().init_some_extra_state(model) # sets position_ids, b_q_seq_len, b_q_start_loc (prefill)
+ pos = self.position_ids
+ self.position_cos_sliding = torch.index_select(model._cos_cached_sliding, 0, pos) # [T, rope_dim//2]
+ self.position_sin_sliding = torch.index_select(model._sin_cached_sliding, 0, pos)
+ self.position_cos_compress = torch.index_select(model._cos_cached_compress, 0, pos)
+ self.position_sin_compress = torch.index_select(model._sin_cached_compress, 0, pos)
+ # Per-token request id (decode: one token per req; prefill: ragged -> repeat by q-len).
+ # Layer-independent; the swa kernel + build_metadata's c4/c128 readers all reuse it.
+ if self.is_prefill:
+ self.dsv4_sparse_req_idx = torch.repeat_interleave(self.b_req_idx, self.b_q_seq_len.long())
+ self._dsv4_token_to_batch_idx = torch.repeat_interleave(
+ torch.arange(self.b_req_idx.shape[0], device=self.b_req_idx.device),
+ self.b_q_seq_len.long(),
+ output_size=pos.numel(),
+ ).to(torch.int32)
+ else:
+ self.dsv4_sparse_req_idx = self.b_req_idx
+ self._dsv4_token_to_batch_idx = None
+ # Sliding-window indices are layer-independent, so build them once into the model workspace.
+ from lightllm.models.deepseek_v4.triton_kernel.build_swa_index_dsv4 import build_swa_index
+
+ workspace = model.dsv4_workspace
+ self.dsv4_workspace = workspace
+ self.dsv4_swa_indices, self.dsv4_swa_lengths = workspace.swa(self.microbatch_index, pos.numel())
+ self.dsv4_swa_indices, self.dsv4_swa_lengths = build_swa_index(
+ req_idx=self.dsv4_sparse_req_idx,
+ positions=self.position_ids,
+ req_to_token_indexs=self.req_manager.req_to_token_indexs,
+ full_to_swa_indexs=self.mem_manager.full_to_swa_indexs,
+ swa_index=self.dsv4_swa_indices,
+ swa_length=self.dsv4_swa_lengths,
+ )
+ from lightllm.models.deepseek_v4.triton_kernel.build_compress_index_dsv4 import build_compress_index
+
+ cap = workspace.compress_cap(self._dsv4_index_max_kv_seq_len(model), 128)
+ self.dsv4_c128_indices, self.dsv4_c128_lengths = workspace.c128(self.microbatch_index, pos.numel(), cap)
+ build_compress_index(
+ self.dsv4_sparse_req_idx,
+ self.position_ids,
+ self.req_manager.req_to_token_indexs,
+ self.mem_manager.full_to_c128_indexs,
+ 128,
+ self.dsv4_c128_indices,
+ self.dsv4_c128_lengths,
+ )
+ # prefill-cudagraph 桶填充的 HOLD 尾请求的 q 行数。其注意力读 HOLD 槽位(内容被并发写
+ # 竞争,每轮不同),输出必须清零,否则 pad 行 hidden 不确定 -> MoE 路由抖动 -> 共享 expert
+ # 批次组成变化 -> 真实行 GEMM 归约顺序变化(ulp 级),44 层放大后翻转低置信 token。
+ self._dsv4_prefill_pad_q_len = 0
+ if self.is_prefill and self.b_req_idx.numel() > 0:
+ if int(self.b_req_idx[-1].item()) == self.req_manager.HOLD_REQUEST_ID:
+ self._dsv4_prefill_pad_q_len = int((self.b_seq_len[-1] - self.b_ready_cache_len[-1]).item())
diff --git a/lightllm/models/deepseek_v4/layer_infer/__init__.py b/lightllm/models/deepseek_v4/layer_infer/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/lightllm/models/deepseek_v4/layer_infer/compressor.py b/lightllm/models/deepseek_v4/layer_infer/compressor.py
new file mode 100644
index 0000000000..e9f2af7385
--- /dev/null
+++ b/lightllm/models/deepseek_v4/layer_infer/compressor.py
@@ -0,0 +1,477 @@
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+from triton.language.extra import libdevice
+
+from lightllm.common.kv_cache_mem_manager import DeepseekV4MemoryManager
+from lightllm.common.kv_cache_mem_manager.deepseek4_mem_manager import (
+ DSV4_C4_STATE_RING,
+ DSV4_C128_STATE_RING,
+ DSV4_SWA_PAGE_SIZE,
+)
+
+
+@dataclass
+class CoreCompressorMetadata:
+ layer_idx: int
+ compress_ratio: int
+ out_slots: torch.Tensor
+ mem_index: torch.Tensor
+ state_buffer: torch.Tensor
+ out_buffer: torch.Tensor
+ out_page_size: int
+ position_ids: torch.Tensor
+ b_req_idx: torch.Tensor
+ b_seq_len: torch.Tensor
+ b_ready_cache_len: Optional[torch.Tensor]
+ b_q_start_loc: Optional[torch.Tensor]
+ req_to_token_indexs: torch.Tensor
+ full_to_swa_indexs: torch.Tensor
+ token_to_batch_idx: Optional[torch.Tensor]
+ kv_score: Optional[torch.Tensor]
+ is_prefill: bool
+
+
+@triton.jit
+def _add_ape_to_kv_score_kernel(
+ kv_score,
+ kv_score_stride0,
+ kv_score_stride1,
+ ape,
+ ape_stride0,
+ positions,
+ STATE_WIDTH: tl.constexpr,
+ COMPRESS_RATIO: tl.constexpr,
+ BLOCK: tl.constexpr,
+):
+ token_idx = tl.program_id(0)
+ offs = tl.arange(0, BLOCK)
+ mask = offs < STATE_WIDTH
+
+ position = tl.load(positions + token_idx)
+ ape_row = position % COMPRESS_RATIO
+ score = tl.load(kv_score + token_idx * kv_score_stride0 + (STATE_WIDTH + offs) * kv_score_stride1, mask=mask)
+ ape_value = tl.load(ape + ape_row * ape_stride0 + offs, mask=mask)
+ tl.store(
+ kv_score + token_idx * kv_score_stride0 + (STATE_WIDTH + offs) * kv_score_stride1,
+ score + ape_value,
+ mask=mask,
+ )
+ return
+
+
+@triton.jit
+def _save_partial_states_kernel(
+ kv_score,
+ kv_score_stride0,
+ kv_score_stride1,
+ positions,
+ token_to_batch_idx,
+ b_req_idx,
+ b_seq_len,
+ mem_index,
+ full_to_swa,
+ state_buffer,
+ STATE_WIDTH: tl.constexpr,
+ STATE_LAST_DIM: tl.constexpr,
+ COMPRESS_RATIO: tl.constexpr,
+ IS_C4: tl.constexpr,
+ IS_PREFILL: tl.constexpr,
+ SWA_PAGE_SIZE: tl.constexpr,
+ STATE_RING: tl.constexpr,
+ BLOCK: tl.constexpr,
+):
+ token_idx = tl.program_id(0)
+ batch_idx = tl.load(token_to_batch_idx + token_idx) if IS_PREFILL else token_idx
+ position = tl.load(positions + token_idx)
+ seq_len = tl.load(b_seq_len + batch_idx)
+
+ if IS_C4:
+ same_page_next = (position % SWA_PAGE_SIZE) + STATE_RING < SWA_PAGE_SIZE
+ if same_page_next and position + STATE_RING < seq_len:
+ return
+ else:
+ if position + COMPRESS_RATIO < seq_len:
+ return
+
+ full_slot = tl.load(mem_index + token_idx).to(tl.int64)
+ swa_slot = tl.load(full_to_swa + full_slot).to(tl.int64)
+ if swa_slot < 0:
+ return
+ state_row = (swa_slot // SWA_PAGE_SIZE) * STATE_RING + (swa_slot % STATE_RING)
+
+ offs = tl.arange(0, BLOCK)
+ mask = offs < STATE_WIDTH
+ kv = tl.load(kv_score + token_idx * kv_score_stride0 + offs * kv_score_stride1, mask=mask)
+ score = tl.load(kv_score + token_idx * kv_score_stride0 + (STATE_WIDTH + offs) * kv_score_stride1, mask=mask)
+ state_base = state_buffer + state_row * STATE_LAST_DIM
+ tl.store(state_base + offs, kv, mask=mask)
+ tl.store(state_base + STATE_WIDTH + offs, score, mask=mask)
+ return
+
+
+@triton.jit
+def _fused_compress_norm_rope_insert_kernel(
+ kv_score,
+ kv_score_stride0,
+ kv_score_stride1,
+ state_buffer,
+ positions,
+ token_to_batch_idx,
+ b_req_idx,
+ b_seq_len,
+ b_ready_cache_len,
+ b_q_start_loc,
+ req_to_token,
+ req_to_token_stride0,
+ full_to_swa,
+ out_slots,
+ norm_weight,
+ rms_eps,
+ cos_table,
+ cos_stride0,
+ cos_stride1,
+ sin_table,
+ sin_stride0,
+ sin_stride1,
+ out_buffer,
+ HEAD_DIM: tl.constexpr,
+ STATE_WIDTH: tl.constexpr,
+ STATE_LAST_DIM: tl.constexpr,
+ COMPRESS_RATIO: tl.constexpr,
+ WINDOW_SIZE: tl.constexpr,
+ IS_C4: tl.constexpr,
+ IS_PREFILL: tl.constexpr,
+ SWA_PAGE_SIZE: tl.constexpr,
+ STATE_RING: tl.constexpr,
+ ROPE_HEAD_DIM: tl.constexpr,
+ FP8_MAX: tl.constexpr,
+ SCALE_MIN: tl.constexpr,
+ NOPE_DIM: tl.constexpr,
+ QUANT_BLOCK: tl.constexpr,
+ SCALE_BYTES: tl.constexpr,
+ PAGE_SIZE: tl.constexpr,
+ BYTES_PER_PAGE: tl.constexpr,
+ BLOCK: tl.constexpr,
+ OUTPUT_BF16: tl.constexpr,
+):
+ token_idx = tl.program_id(0)
+ out_slot = tl.load(out_slots + token_idx).to(tl.int64)
+ if out_slot < 0:
+ return
+
+ position = tl.load(positions + token_idx)
+ if (position + 1) % COMPRESS_RATIO != 0:
+ return
+
+ batch_idx = tl.load(token_to_batch_idx + token_idx) if IS_PREFILL else token_idx
+ req_idx = tl.load(b_req_idx + batch_idx).to(tl.int64)
+ seq_len = tl.load(b_seq_len + batch_idx)
+ if IS_PREFILL:
+ ready_len = tl.load(b_ready_cache_len + batch_idx)
+ q_start = tl.load(b_q_start_loc + batch_idx)
+ else:
+ ready_len = position
+ q_start = token_idx
+
+ token_offsets = tl.arange(0, WINDOW_SIZE)
+ start = position - WINDOW_SIZE + 1
+ gather_pos = start + token_offsets
+ valid_pos = (gather_pos >= 0) & (gather_pos < seq_len)
+ use_current = (gather_pos >= ready_len) & valid_pos if IS_PREFILL else gather_pos == position
+ current_idx = q_start + (gather_pos - ready_len) if IS_PREFILL else token_idx + token_offsets * 0
+
+ if IS_C4:
+ full_slot = tl.load(
+ req_to_token + req_idx * req_to_token_stride0 + gather_pos,
+ mask=valid_pos & (~use_current),
+ other=0,
+ ).to(tl.int64)
+ swa_slot = tl.load(full_to_swa + full_slot, mask=valid_pos & (~use_current), other=-1).to(tl.int64)
+ state_row = (swa_slot // SWA_PAGE_SIZE) * STATE_RING + (swa_slot % STATE_RING)
+ state_valid = valid_pos & (~use_current) & (swa_slot >= 0)
+ head_offset = tl.where(token_offsets >= COMPRESS_RATIO, HEAD_DIM, 0)
+ else:
+ full_slot = tl.load(
+ req_to_token + req_idx * req_to_token_stride0 + gather_pos,
+ mask=valid_pos & (~use_current),
+ other=0,
+ ).to(tl.int64)
+ swa_slot = tl.load(full_to_swa + full_slot, mask=valid_pos & (~use_current), other=-1).to(tl.int64)
+ state_row = (swa_slot // SWA_PAGE_SIZE) * STATE_RING + (swa_slot % STATE_RING)
+ state_valid = valid_pos & (~use_current) & (swa_slot >= 0)
+ head_offset = token_offsets * 0
+
+ offs = tl.arange(0, BLOCK)
+ dim_mask = offs < HEAD_DIM
+ current_mask = use_current[:, None] & dim_mask[None, :]
+ state_mask = state_valid[:, None] & dim_mask[None, :]
+
+ cur_kv = tl.load(
+ kv_score + current_idx[:, None] * kv_score_stride0 + (head_offset[:, None] + offs[None, :]) * kv_score_stride1,
+ mask=current_mask,
+ other=0.0,
+ )
+ cur_score = tl.load(
+ kv_score
+ + current_idx[:, None] * kv_score_stride0
+ + (STATE_WIDTH + head_offset[:, None] + offs[None, :]) * kv_score_stride1,
+ mask=current_mask,
+ other=float("-inf"),
+ )
+ state_kv = tl.load(
+ state_buffer + state_row[:, None] * STATE_LAST_DIM + head_offset[:, None] + offs[None, :],
+ mask=state_mask,
+ other=0.0,
+ )
+ state_score = tl.load(
+ state_buffer + state_row[:, None] * STATE_LAST_DIM + STATE_WIDTH + head_offset[:, None] + offs[None, :],
+ mask=state_mask,
+ other=float("-inf"),
+ )
+
+ kv = tl.where(current_mask, cur_kv, state_kv)
+ score = tl.where(current_mask, cur_score, state_score)
+ score = tl.softmax(score, dim=0)
+ compressed_kv = tl.sum(kv * score, axis=0)
+
+ rms_w = tl.load(norm_weight + offs, mask=dim_mask, other=0.0)
+ variance = tl.sum(compressed_kv * compressed_kv, axis=0) / HEAD_DIM
+ rrms = tl.rsqrt(variance + rms_eps)
+ normed = compressed_kv * rrms * rms_w
+
+ num_pairs: tl.constexpr = BLOCK // 2
+ nope_pairs: tl.constexpr = NOPE_DIM // 2
+ pair_2d = tl.reshape(normed, (num_pairs, 2))
+ even, odd = tl.split(pair_2d)
+ pair_idx = tl.arange(0, num_pairs)
+ rope_pair_local = pair_idx - nope_pairs
+ is_rope_pair = rope_pair_local >= 0
+ cs_idx = tl.maximum(rope_pair_local, 0)
+ compressed_pos = (position // COMPRESS_RATIO) * COMPRESS_RATIO
+ cos_v = tl.load(cos_table + compressed_pos * cos_stride0 + cs_idx * cos_stride1, mask=is_rope_pair, other=1.0)
+ sin_v = tl.load(sin_table + compressed_pos * sin_stride0 + cs_idx * sin_stride1, mask=is_rope_pair, other=0.0)
+ new_even = even * cos_v - odd * sin_v
+ new_odd = odd * cos_v + even * sin_v
+ rotated = tl.interleave(new_even, new_odd)
+
+ if OUTPUT_BF16:
+ # indexer-K path: emit the post-rope full HEAD_DIM vector as dense bf16 (token-indexed),
+ # leaving the fp8 single-amax pack to destindex_copy_indexer_k_dsv4 (the c4_indexer_pool
+ # ABI differs from the latent slab: whole-vector fp8 + one fp32 scale, no bf16 rope tail).
+ tl.store(out_buffer + token_idx * HEAD_DIM + offs, rotated.to(tl.bfloat16), mask=dim_mask)
+ return
+
+ page = out_slot // PAGE_SIZE
+ token_in_page = out_slot % PAGE_SIZE
+ data_base = page * BYTES_PER_PAGE + token_in_page * (NOPE_DIM + ROPE_HEAD_DIM * 2)
+ scale_base = page * BYTES_PER_PAGE + PAGE_SIZE * (NOPE_DIM + ROPE_HEAD_DIM * 2) + token_in_page * SCALE_BYTES
+
+ n_quant_blocks: tl.constexpr = BLOCK // QUANT_BLOCK
+ n_nope_blocks: tl.constexpr = NOPE_DIM // QUANT_BLOCK
+ quant_input = normed.to(tl.bfloat16).to(tl.float32)
+ quant_2d = tl.reshape(quant_input, (n_quant_blocks, QUANT_BLOCK))
+ abs_2d = tl.abs(quant_2d)
+ block_absmax = tl.max(abs_2d, axis=1)
+ scale_exp = tl.ceil(libdevice.log2(tl.maximum(block_absmax / FP8_MAX, SCALE_MIN))).to(tl.int32)
+ scale = ((scale_exp + 127) << 23).to(tl.float32, bitcast=True)
+ kv_fp8 = tl.clamp(quant_2d / scale[:, None], -FP8_MAX, FP8_MAX).to(tl.float8e4nv)
+ kv_u8 = tl.reshape(kv_fp8.to(tl.uint8, bitcast=True), (BLOCK,))
+ tl.store(out_buffer + data_base + offs, kv_u8, mask=offs < NOPE_DIM)
+
+ scale_idx = tl.arange(0, SCALE_BYTES)
+ scale_bytes = tl.where(scale_idx < n_nope_blocks, scale_exp + 127, 0).to(tl.uint8)
+ tl.store(out_buffer + scale_base + scale_idx, scale_bytes)
+
+ rope_local = offs - NOPE_DIM
+ rope_mask = (offs >= NOPE_DIM) & dim_mask
+ rope_ptr = (out_buffer + data_base + NOPE_DIM).to(tl.pointer_type(tl.bfloat16))
+ tl.store(rope_ptr + rope_local, rotated.to(tl.bfloat16), mask=rope_mask)
+ return
+
+
+def prepare_compress_states(*, infer_state, layer_idx: int, compress_ratio: int, is_in_indexer: bool = False):
+ if compress_ratio == 0:
+ return None
+
+ mem_manager: DeepseekV4MemoryManager = infer_state.mem_manager
+ if is_in_indexer:
+ # c4 Lightning-Indexer key compression: same window/state machinery as the c4 latent
+ # compressor but with index_head_dim, a separate state pool, and a DENSE bf16 scratch
+ # out_buffer (the kernel's OUTPUT_BF16 path); the fp8 pack into c4_indexer_pool is done
+ # afterwards by pack_indexer_k_to_cache.
+ assert compress_ratio == 4, "只有 c4(CSA) 层有 indexer-K"
+ out_slots = mem_manager.full_to_c4_indexs[infer_state.mem_index.long().reshape(-1)]
+ state_buffer = mem_manager.get_c4_indexer_state_buffer(layer_idx)
+ out_buffer = torch.empty(
+ (infer_state.mem_index.numel(), mem_manager.indexer_head_dim),
+ dtype=torch.bfloat16,
+ device=infer_state.mem_index.device,
+ )
+ out_page_size = 1 # unused under OUTPUT_BF16 (token-indexed dense scratch, not paged)
+ else:
+ if compress_ratio == 4:
+ out_slots = mem_manager.full_to_c4_indexs[infer_state.mem_index.long().reshape(-1)]
+ state_buffer = mem_manager.get_c4_state_buffer(layer_idx)
+ out_pool = mem_manager.c4_pool
+ elif compress_ratio == 128:
+ out_slots = mem_manager.full_to_c128_indexs[infer_state.mem_index.long().reshape(-1)]
+ state_buffer = mem_manager.get_c128_state_buffer(layer_idx)
+ out_pool = mem_manager.c128_pool
+ else:
+ raise AssertionError(f"invalid DeepSeek-V4 compress ratio {compress_ratio}")
+ out_buffer = mem_manager.get_compressed_kv_buffer(layer_idx)
+ out_page_size = out_pool.page_size
+
+ token_to_batch_idx = infer_state.b_req_idx
+ if infer_state.is_prefill:
+ token_to_batch_idx = getattr(infer_state, "_dsv4_token_to_batch_idx", None)
+ if token_to_batch_idx is None or token_to_batch_idx.numel() != infer_state.position_ids.numel():
+ q_lens = (infer_state.b_seq_len - infer_state.b_ready_cache_len).to(torch.long)
+ batch_idx = torch.arange(infer_state.b_req_idx.shape[0], device=infer_state.b_req_idx.device)
+ token_to_batch_idx = torch.repeat_interleave(
+ batch_idx, q_lens, output_size=infer_state.position_ids.numel()
+ ).to(torch.int32)
+ infer_state._dsv4_token_to_batch_idx = token_to_batch_idx
+
+ return CoreCompressorMetadata(
+ layer_idx=layer_idx,
+ compress_ratio=compress_ratio,
+ out_slots=out_slots,
+ mem_index=infer_state.mem_index,
+ state_buffer=state_buffer,
+ out_buffer=out_buffer,
+ out_page_size=out_page_size,
+ position_ids=infer_state.position_ids,
+ b_req_idx=infer_state.b_req_idx,
+ b_seq_len=infer_state.b_seq_len,
+ b_ready_cache_len=infer_state.b_ready_cache_len,
+ b_q_start_loc=infer_state.b_q_start_loc,
+ req_to_token_indexs=infer_state.req_manager.req_to_token_indexs,
+ full_to_swa_indexs=mem_manager.full_to_swa_indexs,
+ token_to_batch_idx=token_to_batch_idx,
+ kv_score=None,
+ is_prefill=infer_state.is_prefill,
+ )
+
+
+def prepare_partial_states(
+ *,
+ kv_score: torch.Tensor,
+ metadata: Optional[CoreCompressorMetadata],
+ ape: torch.Tensor,
+ compress_ratio: int,
+):
+ if metadata is None or kv_score.shape[0] == 0:
+ return
+ state_width = kv_score.shape[-1] // 2
+ _add_ape_to_kv_score_kernel[(kv_score.shape[0],)](
+ kv_score,
+ kv_score.stride(0),
+ kv_score.stride(1),
+ ape,
+ ape.stride(0),
+ metadata.position_ids,
+ STATE_WIDTH=state_width,
+ COMPRESS_RATIO=compress_ratio,
+ BLOCK=triton.next_power_of_2(state_width),
+ num_warps=4,
+ )
+ return
+
+
+def fused_compress(
+ *,
+ kv_score: torch.Tensor,
+ metadata: Optional[CoreCompressorMetadata],
+ norm_weight: torch.Tensor,
+ ape: torch.Tensor,
+ eps: float,
+ head_dim: int,
+ qk_rope_head_dim: int,
+ compress_ratio: int,
+ cos_table: torch.Tensor,
+ sin_table: torch.Tensor,
+ output_bf16: bool = False,
+):
+ if metadata is None or kv_score.shape[0] == 0:
+ return
+
+ state_width = kv_score.shape[-1] // 2
+ state_last_dim = metadata.state_buffer.shape[-1]
+ is_c4 = compress_ratio == 4
+ state_ring = DSV4_C4_STATE_RING if is_c4 else DSV4_C128_STATE_RING
+ block_state = triton.next_power_of_2(state_width)
+ block_head = triton.next_power_of_2(head_dim)
+
+ _fused_compress_norm_rope_insert_kernel[(kv_score.shape[0],)](
+ kv_score,
+ kv_score.stride(0),
+ kv_score.stride(1),
+ metadata.state_buffer,
+ metadata.position_ids,
+ metadata.token_to_batch_idx,
+ metadata.b_req_idx,
+ metadata.b_seq_len,
+ metadata.b_ready_cache_len if metadata.b_ready_cache_len is not None else metadata.b_seq_len,
+ metadata.b_q_start_loc if metadata.b_q_start_loc is not None else metadata.b_seq_len,
+ metadata.req_to_token_indexs,
+ metadata.req_to_token_indexs.stride(0),
+ metadata.full_to_swa_indexs,
+ metadata.out_slots,
+ norm_weight,
+ eps,
+ cos_table,
+ cos_table.stride(0),
+ cos_table.stride(1),
+ sin_table,
+ sin_table.stride(0),
+ sin_table.stride(1),
+ metadata.out_buffer,
+ HEAD_DIM=head_dim,
+ STATE_WIDTH=state_width,
+ STATE_LAST_DIM=state_last_dim,
+ COMPRESS_RATIO=compress_ratio,
+ WINDOW_SIZE=compress_ratio * (2 if is_c4 else 1),
+ IS_C4=is_c4,
+ IS_PREFILL=metadata.is_prefill,
+ SWA_PAGE_SIZE=DSV4_SWA_PAGE_SIZE,
+ STATE_RING=state_ring,
+ ROPE_HEAD_DIM=qk_rope_head_dim,
+ FP8_MAX=torch.finfo(torch.float8_e4m3fn).max,
+ SCALE_MIN=1e-4,
+ NOPE_DIM=head_dim - qk_rope_head_dim,
+ QUANT_BLOCK=64,
+ SCALE_BYTES=(head_dim - qk_rope_head_dim) // 64 + 1,
+ PAGE_SIZE=metadata.out_page_size,
+ BYTES_PER_PAGE=metadata.out_buffer.shape[-1],
+ BLOCK=block_head,
+ OUTPUT_BF16=output_bf16,
+ num_warps=4,
+ )
+
+ _save_partial_states_kernel[(kv_score.shape[0],)](
+ kv_score,
+ kv_score.stride(0),
+ kv_score.stride(1),
+ metadata.position_ids,
+ metadata.token_to_batch_idx,
+ metadata.b_req_idx,
+ metadata.b_seq_len,
+ metadata.mem_index,
+ metadata.full_to_swa_indexs,
+ metadata.state_buffer,
+ STATE_WIDTH=state_width,
+ STATE_LAST_DIM=state_last_dim,
+ COMPRESS_RATIO=compress_ratio,
+ IS_C4=is_c4,
+ IS_PREFILL=metadata.is_prefill,
+ SWA_PAGE_SIZE=DSV4_SWA_PAGE_SIZE,
+ STATE_RING=state_ring,
+ BLOCK=block_state,
+ num_warps=4,
+ )
+ return
diff --git a/lightllm/models/deepseek_v4/layer_infer/hyper_connection.py b/lightllm/models/deepseek_v4/layer_infer/hyper_connection.py
new file mode 100644
index 0000000000..080ebabd89
--- /dev/null
+++ b/lightllm/models/deepseek_v4/layer_infer/hyper_connection.py
@@ -0,0 +1,75 @@
+import torch
+
+try:
+ import vllm.model_executor.layers.mhc # noqa: F401
+except Exception as e:
+ raise RuntimeError("DeepSeek-V4 requires vLLM mHC custom ops; failed to import vllm MHC kernels") from e
+
+
+# vllm DeepseekV4DecoderLayer.hc_post_alpha
+HC_POST_ALPHA = 2.0
+
+
+def hc_pre(residual, hc_fn, hc_scale, hc_base, rms_eps, hc_eps, sinkhorn_iters, norm_weight, norm_eps):
+ """Standalone hc_pre for the first layer. residual:[T, hc, dim] ->
+ (x[T,dim], residual, post_mix[T,hc,1], res_mix[T,hc,hc]); the sub-layer RMSNorm is fused via norm_weight."""
+ post_mix, res_mix, x = torch.ops.vllm.mhc_pre_tilelang(
+ residual=residual,
+ fn=hc_fn,
+ hc_scale=hc_scale,
+ hc_base=hc_base,
+ rms_eps=rms_eps,
+ hc_pre_eps=hc_eps,
+ hc_sinkhorn_eps=hc_eps,
+ hc_post_mult_value=HC_POST_ALPHA,
+ sinkhorn_repeat=sinkhorn_iters,
+ norm_weight=norm_weight,
+ norm_eps=norm_eps,
+ )
+ return x, residual, post_mix, res_mix
+
+
+def hc_fused_post_pre(
+ x, residual, post_mix, res_mix, hc_fn, hc_scale, hc_base, rms_eps, hc_eps, sinkhorn_iters, norm_weight, norm_eps
+):
+ """hc_post of the previous sub-layer fused with hc_pre of the next one (norm fused too).
+ Returns (x[T,dim], residual[T,hc,dim], post_mix, res_mix)."""
+ residual, post_mix, res_mix, x = torch.ops.vllm.mhc_fused_post_pre_tilelang(
+ x=x,
+ residual=residual,
+ post_layer_mix=post_mix,
+ comb_res_mix=res_mix,
+ fn=hc_fn,
+ hc_scale=hc_scale,
+ hc_base=hc_base,
+ rms_eps=rms_eps,
+ hc_pre_eps=hc_eps,
+ hc_sinkhorn_eps=hc_eps,
+ hc_post_mult_value=HC_POST_ALPHA,
+ sinkhorn_repeat=sinkhorn_iters,
+ norm_weight=norm_weight,
+ norm_eps=norm_eps,
+ )
+ return x, residual, post_mix, res_mix
+
+
+def hc_post(x, residual, post_mix, res_mix):
+ """Complete the hc_post left pending by the last layer. -> streams [T, hc, dim]."""
+ return torch.ops.vllm.mhc_post_tilelang(x, residual, post_mix, res_mix)
+
+
+def hc_head(streams, hc_fn, hc_scale, hc_base, hc_mult, dim, rms_eps, hc_eps, alloc_func):
+ """Final stream collapse before the lm_head. streams:[N, hc*dim] -> [N, dim]."""
+ out = alloc_func((streams.shape[0], dim), dtype=streams.dtype, device=streams.device)
+ torch.ops.vllm.hc_head_fused_kernel_tilelang(
+ streams.view(-1, hc_mult, dim).contiguous(),
+ hc_fn,
+ hc_scale,
+ hc_base,
+ out,
+ dim,
+ rms_eps,
+ hc_eps,
+ hc_mult,
+ )
+ return out
diff --git a/lightllm/models/deepseek_v4/layer_infer/post_layer_infer.py b/lightllm/models/deepseek_v4/layer_infer/post_layer_infer.py
new file mode 100644
index 0000000000..0eb5cfc2b6
--- /dev/null
+++ b/lightllm/models/deepseek_v4/layer_infer/post_layer_infer.py
@@ -0,0 +1,28 @@
+from lightllm.models.llama.layer_infer.post_layer_infer import LlamaPostLayerInfer
+from .hyper_connection import hc_head, hc_post
+from ..infer_struct import DeepseekV4InferStateInfo
+
+
+class DeepseekV4PostLayerInfer(LlamaPostLayerInfer):
+ """Collapse the hc_mult residual streams (hc_head) to [T, hidden], then final norm + lm_head."""
+
+ def token_forward(self, input_embdings, infer_state: DeepseekV4InferStateInfo, layer_weight):
+ cfg = layer_weight.network_config_
+ if isinstance(input_embdings, tuple):
+ # truncated-layer runs (autotune warmup) end before the last layer's _hc_ffn_out
+ # collapse; finish the pending hc_post here.
+ streams = hc_post(*input_embdings)
+ input_embdings = streams.reshape(streams.shape[0], -1)
+ collapsed = hc_head(
+ input_embdings,
+ layer_weight.hc_head_fn_.weight,
+ layer_weight.hc_head_scale_.weight,
+ layer_weight.hc_head_base_.weight,
+ cfg["hc_mult"],
+ cfg["hidden_size"],
+ cfg["rms_norm_eps"],
+ cfg.get("hc_eps", 1e-6),
+ self.alloc_tensor,
+ )
+ logits = super().token_forward(collapsed, infer_state, layer_weight)
+ return logits, input_embdings
diff --git a/lightllm/models/deepseek_v4/layer_infer/pre_layer_infer.py b/lightllm/models/deepseek_v4/layer_infer/pre_layer_infer.py
new file mode 100644
index 0000000000..b95f5a14a8
--- /dev/null
+++ b/lightllm/models/deepseek_v4/layer_infer/pre_layer_infer.py
@@ -0,0 +1,24 @@
+import torch
+import torch.distributed as dist
+from lightllm.models.llama.layer_infer.pre_layer_infer import LlamaPreLayerInfer
+from lightllm.distributed.communication_op import all_reduce
+from ..infer_struct import DeepseekV4InferStateInfo
+
+
+class DeepseekV4PreLayerInfer(LlamaPreLayerInfer):
+ """Token embedding, then expand to the hc_mult parallel residual streams [T, hc_mult*hidden]."""
+
+ def __init__(self, network_config):
+ super().__init__(network_config)
+ self.hc_mult = network_config["hc_mult"]
+ return
+
+ def context_forward(self, input_ids, infer_state: DeepseekV4InferStateInfo, layer_weight):
+ input_embdings = super().context_forward(input_ids, infer_state, layer_weight)
+ t, hidden = input_embdings.shape
+ return input_embdings.unsqueeze(1).expand(t, self.hc_mult, hidden).reshape(t, self.hc_mult * hidden)
+
+ def token_forward(self, input_ids, infer_state: DeepseekV4InferStateInfo, layer_weight):
+ input_embdings = super().token_forward(input_ids, infer_state, layer_weight)
+ t, hidden = input_embdings.shape
+ return input_embdings.unsqueeze(1).expand(t, self.hc_mult, hidden).reshape(t, self.hc_mult * hidden)
diff --git a/lightllm/models/deepseek_v4/layer_infer/transformer_layer_infer.py b/lightllm/models/deepseek_v4/layer_infer/transformer_layer_infer.py
new file mode 100644
index 0000000000..638e11f6f2
--- /dev/null
+++ b/lightllm/models/deepseek_v4/layer_infer/transformer_layer_infer.py
@@ -0,0 +1,803 @@
+import torch
+import torch.distributed as dist
+from lightllm.common.basemodel import TransformerLayerInferTpl
+from lightllm.common.basemodel.attention.base_att import AttControl
+from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import silu_and_mul_fwd
+from lightllm.distributed.communication_op import all_reduce
+from lightllm.models.deepseek3_2.layer_infer.transformer_layer_infer import Deepseek3_2TransformerLayerInfer
+from lightllm.models.deepseek_v4.layer_weights.transformer_layer_weight import DeepseekV4TransformerLayerWeight
+from lightllm.utils.envs_utils import get_env_start_args
+from lightllm.utils.tensor_utils import tensor_to_no_ref_tensor
+from lightllm.utils.vllm_utils import vllm_ops
+from .hyper_connection import hc_pre, hc_fused_post_pre, hc_post
+from .compressor import fused_compress as fused_compress_op
+from .compressor import prepare_partial_states
+from .compressor import prepare_compress_states
+from lightllm.models.deepseek2.triton_kernel.rotary_emb import rotary_emb_fwd
+from ..infer_struct import DeepseekV4InferStateInfo
+import deep_gemm
+from lightllm.models.deepseek_v4.triton_kernel.topk_transform import topk_transform_512
+
+
+_C4_PREFILL_LOGITS_BUDGET_BYTES = 512 * 1024 * 1024
+
+
+class DeepseekV4TransformerLayerInfer(Deepseek3_2TransformerLayerInfer):
+ def __init__(self, layer_num, network_config):
+ TransformerLayerInferTpl.__init__(self, layer_num, network_config)
+ self.eps_ = network_config["rms_norm_eps"]
+ self.embed_dim_ = network_config["hidden_size"]
+ self.num_heads = network_config["num_attention_heads"]
+ self.head_dim_ = network_config["head_dim"]
+ self.qk_rope_head_dim = network_config["qk_rope_head_dim"]
+ self.qk_nope_head_dim = self.head_dim_ - self.qk_rope_head_dim
+ self.v_head_dim = self.head_dim_
+ self.o_groups = network_config["o_groups"]
+ self.hc_mult = network_config["hc_mult"]
+ self.sinkhorn_iters = network_config["hc_sinkhorn_iters"]
+ self.hc_eps = network_config["hc_eps"]
+ self.compress_ratio = network_config["compress_ratios"][layer_num]
+ self.is_hash = layer_num < network_config["num_hash_layers"]
+ self.is_last_layer = layer_num == network_config["n_layer"] - 1
+ # complex64 rope table for this layer's variant (sliding / compressed); set by
+ # DeepseekV4TpPartModel._init_to_get_rotary once the tables are built. The full compress
+ # cos/sin tables (compressor entry rope uses entry positions, not token positions) are
+ # wired there too.
+ self.freqs_cis = None
+ self.cos_compress_table = None
+ self.sin_compress_table = None
+ self.num_experts_per_tok = network_config["num_experts_per_tok"]
+ self.routed_scaling_factor = network_config["routed_scaling_factor"]
+ self.swiglu_limit = float(network_config["swiglu_limit"])
+ self.softmax_scale = (self.qk_nope_head_dim + self.qk_rope_head_dim) ** (-0.5)
+ self.tp_q_head_num_ = self.num_heads // self.tp_world_size_
+ self.flashmla_q_head_num_ = self.tp_q_head_num_
+ self.tp_groups = self.o_groups // self.tp_world_size_
+ self.enable_ep_moe = get_env_start_args().enable_ep_moe
+ self.compressor = CompressorInfer(
+ layer_idx=self.layer_num_, network_config=self.network_config_, tp_world_size=self.tp_world_size_
+ )
+ self.index_infer = DeepseekV4IndexInfer(
+ layer_idx=self.layer_num_, network_config=self.network_config_, tp_world_size=self.tp_world_size_
+ )
+ self.dsv4_prefill_aux_stream = None
+
+ # ------------------------------------------------------------------ forward (HC-threaded)
+ def _hc_attn_in(self, input_embdings, layer_weight: DeepseekV4TransformerLayerWeight):
+ """Layer input -> attention input (attn_norm fused). First layer gets the raw streams
+ and runs a standalone hc_pre; later layers get (x, residual, post_mix, res_mix) and fuse
+ the previous layer's ffn hc_post with this layer's attn hc_pre."""
+ if torch.is_tensor(input_embdings):
+ residual = input_embdings.view(-1, self.hc_mult, self.embed_dim_)
+ return hc_pre(
+ residual,
+ layer_weight.hc_attn_fn_.weight,
+ layer_weight.hc_attn_scale_.weight,
+ layer_weight.hc_attn_base_.weight,
+ self.eps_,
+ self.hc_eps,
+ self.sinkhorn_iters,
+ layer_weight.attn_norm_.weight,
+ self.eps_,
+ )
+ x, residual, post_mix, res_mix = input_embdings
+ return hc_fused_post_pre(
+ x,
+ residual,
+ post_mix,
+ res_mix,
+ layer_weight.hc_attn_fn_.weight,
+ layer_weight.hc_attn_scale_.weight,
+ layer_weight.hc_attn_base_.weight,
+ self.eps_,
+ self.hc_eps,
+ self.sinkhorn_iters,
+ layer_weight.attn_norm_.weight,
+ self.eps_,
+ )
+
+ def _hc_ffn_in(self, x, residual, post_mix, res_mix, layer_weight: DeepseekV4TransformerLayerWeight):
+ """Attention output -> ffn input (ffn_norm fused): fused attn hc_post + ffn hc_pre."""
+ return hc_fused_post_pre(
+ x,
+ residual,
+ post_mix,
+ res_mix,
+ layer_weight.hc_ffn_fn_.weight,
+ layer_weight.hc_ffn_scale_.weight,
+ layer_weight.hc_ffn_base_.weight,
+ self.eps_,
+ self.hc_eps,
+ self.sinkhorn_iters,
+ layer_weight.ffn_norm_.weight,
+ self.eps_,
+ )
+
+ def _hc_ffn_out(self, x, residual, post_mix, res_mix):
+ """Mid layers leave the ffn hc_post pending for the next layer's fused post+pre; the last
+ layer completes it and hands the flat streams [T, hc_mult*hidden] back to the model loop."""
+ if not self.is_last_layer:
+ return x, residual, post_mix, res_mix
+ streams = hc_post(x, residual, post_mix, res_mix)
+ return streams.reshape(streams.shape[0], -1)
+
+ def context_forward(
+ self, input_embdings, infer_state: DeepseekV4InferStateInfo, layer_weight: DeepseekV4TransformerLayerWeight
+ ):
+ x, residual, post_mix, res_mix = self._hc_attn_in(input_embdings, layer_weight)
+ x = self.context_attention_forward(x, infer_state, layer_weight)
+ x, residual, post_mix, res_mix = self._hc_ffn_in(x, residual, post_mix, res_mix, layer_weight)
+ x = self._ffn(x, infer_state, layer_weight)
+ out = self._hc_ffn_out(x, residual, post_mix, res_mix)
+ return out
+
+ def token_forward(
+ self, input_embdings, infer_state: DeepseekV4InferStateInfo, layer_weight: DeepseekV4TransformerLayerWeight
+ ):
+ x, residual, post_mix, res_mix = self._hc_attn_in(input_embdings, layer_weight)
+ x = self.token_attention_forward(x, infer_state, layer_weight)
+ x, residual, post_mix, res_mix = self._hc_ffn_in(x, residual, post_mix, res_mix, layer_weight)
+ x = self._ffn(x, infer_state, layer_weight)
+ return self._hc_ffn_out(x, residual, post_mix, res_mix)
+
+ # ------------------------------------------------------------------ shared projections / cache
+ def _select_rope(self, infer_state: DeepseekV4InferStateInfo):
+ if self.compress_ratio:
+ return infer_state.position_cos_compress, infer_state.position_sin_compress
+ return infer_state.position_cos_sliding, infer_state.position_sin_sliding
+
+ def _get_qkv(
+ self,
+ input: torch.Tensor,
+ infer_state: DeepseekV4InferStateInfo,
+ layer_weight: DeepseekV4TransformerLayerWeight,
+ ):
+ from lightllm.models.deepseek_v4.triton_kernel.norm_rope_cuda import fused_q_norm_rope
+
+ input = self._tpsp_allgather(input=input, infer_state=infer_state)
+ T = input.shape[0]
+ # wq_a and wkv share `input` -> one fused fp8 GEMM, split [q_lora_rank | head_dim]. qa is a
+ # row-strided view (rmsnorm honors stride(0)); kv feeds the fused cache writer -> contiguous.
+ qkv = layer_weight.wq_a_wkv_.mm(input)
+ qa = layer_weight.q_norm_(qkv[:, : -self.head_dim_], eps=self.eps_)
+ q_in = layer_weight.wq_b_.mm(qa).view(T, self.tp_q_head_num_, self.head_dim_)
+ # per-(token, head) weightless self-RMSNorm + interleaved rope on the last rope_dim dims,
+ # fused in one DSV4 CUDA kernel (fp32 norm/rotation, bf16 in between -- same as eager).
+ # The selected FlashMLA MODEL1 binary only instantiates H=64/128. Produce that ABI layout
+ # directly. Prefill uses one max-token workspace so changing T only changes the prefix view;
+ # decode keeps its graph-owned tensor. The workspace's padded tail is zeroed once at init.
+ if infer_state.is_prefill:
+ q = infer_state.dsv4_workspace.flashmla_prefill_q[:T]
+ else:
+ q = self.alloc_tensor((T, self.flashmla_q_head_num_, self.head_dim_), dtype=q_in.dtype, device=q_in.device)
+ q[:, self.tp_q_head_num_ :, :].zero_()
+ fused_q_norm_rope(q_in, q[:, : self.tp_q_head_num_, :], self.eps_, self.freqs_cis, infer_state.position_ids)
+ # kv: rmsnorm + rope + fp8 pack + scatter 进 swa 池,一个 DSV4 CUDA kernel 完成，
+ # 替代 eager norm/rope/cat + _post_cache_kv。
+ # bf16 kv 中间量没有其他消费者: flashmla 路径注意力读 cache,压缩器/indexer 取 x。
+ infer_state.mem_manager.pack_mla_kv_to_cache_fused_norm_rope(
+ layer_index=self.layer_num_,
+ mem_index=infer_state.mem_index,
+ kv=qkv[:, -self.head_dim_ :].contiguous(),
+ kv_weight=layer_weight.kv_norm_.weight,
+ eps=self.eps_,
+ freqs_cis=self.freqs_cis,
+ positions=infer_state.position_ids,
+ )
+ return q, qa, input
+
+ def _get_o(self, o, infer_state: DeepseekV4InferStateInfo, layer_weight: DeepseekV4TransformerLayerWeight):
+ # o: [T, tp_q_head_num_, head_dim_] after inverse rope -> grouped low-rank O -> [T, embed_dim_]
+ position_cos, position_sin = self._select_rope(infer_state)
+ rotary_emb_fwd(o[..., -self.qk_rope_head_dim :], None, position_cos, position_sin, inverse=True)
+ T = o.shape[0]
+ if layer_weight.o_proj_fp8:
+ # one group per rank -> a single fp8 GEMM (deepgemm .mm quantizes o to fp8 internally)
+ o = layer_weight.wo_a_.mm(o.reshape(T, -1)) # [T, o_lora]
+ else:
+ o = o.reshape(T, self.tp_groups, -1).transpose(0, 1).contiguous() # [groups, T, per_group_in]
+ o = layer_weight.wo_a_.bmm(o).transpose(0, 1).reshape(T, -1) # [T, groups*o_lora]
+ o = layer_weight.wo_b_.mm(o)
+ return self._tpsp_reduce(input=o, infer_state=infer_state)
+
+ # ------------------------------------------------------------------ attention (prefill)
+ def context_attention_forward(
+ self, x, infer_state: DeepseekV4InferStateInfo, layer_weight: DeepseekV4TransformerLayerWeight
+ ):
+ # _get_qkv writes the chunk's packed latent into the swa pool (fused kernel) before
+ # attention reads it back via full_to_swa indices (this custom forward bypasses the
+ # tpl _post_cache_kv path).
+ q, q_lora, full_x = self._get_qkv(x, infer_state, layer_weight)
+ o = self._context_attention_wrapper_run(q, q_lora, full_x, infer_state, layer_weight)
+ return self._get_o(o, infer_state, layer_weight)
+
+ def _context_attention_wrapper_run(
+ self, q, q_lora, x, infer_state: DeepseekV4InferStateInfo, layer_weight: DeepseekV4TransformerLayerWeight
+ ):
+ if torch.cuda.is_current_stream_capturing():
+ q = q.contiguous()
+ q_lora = q_lora.contiguous()
+ x = x.contiguous()
+ _q = tensor_to_no_ref_tensor(q)
+ _q_lora = tensor_to_no_ref_tensor(q_lora)
+ _x = tensor_to_no_ref_tensor(x)
+
+ pre_capture_graph = infer_state.prefill_cuda_graph_get_current_capture_graph()
+ pre_capture_graph.__exit__(None, None, None)
+
+ infer_state.prefill_cuda_graph_create_graph_obj()
+ infer_state.prefill_cuda_graph_get_current_capture_graph().__enter__()
+ # Same graph-split output handoff as the template, but avoid its dry-run because
+ # DSV4 attention mutates compressor/cache state before returning.
+ o = self.alloc_tensor((q.shape[0], self.tp_q_head_num_, self.head_dim_), dtype=q.dtype, device=q.device)
+ _o = tensor_to_no_ref_tensor(o)
+
+ def att_func(new_infer_state: DeepseekV4InferStateInfo):
+ self._context_attention_kernel(_q, _q_lora, _x, new_infer_state, layer_weight, out=_o)
+ return
+
+ infer_state.prefill_cuda_graph_add_cpu_runnning_func(func=att_func, after_graph=pre_capture_graph)
+ return o
+
+ return self._context_attention_kernel(q, q_lora, x, infer_state, layer_weight)
+
+ def _compress_and_index(self, q_lora, x, infer_state: DeepseekV4InferStateInfo, layer_weight):
+ cos_table, sin_table = self.cos_compress_table, self.sin_compress_table
+ aux_stream = self.dsv4_prefill_aux_stream
+ if self.compress_ratio == 4 and aux_stream is not None and not torch.cuda.is_current_stream_capturing():
+ # _dsv4_token_to_batch_idx is built in init_some_extra_state (default stream, before this fork),
+ # so both the aux indexer-compressor and the main compressor read a ready, race-free tensor.
+ main_stream = torch.cuda.current_stream()
+ aux_stream.wait_stream(main_stream) # fork: aux waits for x / q_lora produced on main
+ with torch.cuda.stream(aux_stream):
+ # x / q_lora are main-allocated and read here -> record so the allocator won't reuse them.
+ x.record_stream(aux_stream)
+ q_lora.record_stream(aux_stream)
+ self.index_infer.write_indexer_k(
+ x, infer_state, layer_weight, cos_table, sin_table, use_custom_tensor_manager=False
+ )
+ meta = self.index_infer.build_metadata(
+ x, q_lora, infer_state, layer_weight, use_custom_tensor_manager=False
+ )
+ self.compressor.prepare_states(x, infer_state, layer_weight)
+ self.compressor.fused_compress(infer_state, layer_weight, cos_table, sin_table)
+ main_stream.wait_stream(aux_stream) # join before prefill_att reads the indices / latent KV
+ # extra_indices / extra_lengths were allocated on aux -> record on main so they survive until consumed.
+ for _t in (meta.get("extra_indices"), meta.get("extra_lengths")):
+ if _t is not None:
+ _t.record_stream(main_stream)
+ return meta
+
+ # serial fallback -- semantics identical to the original sequence.
+ self.compressor.prepare_states(x, infer_state, layer_weight)
+ self.compressor.fused_compress(infer_state, layer_weight, cos_table, sin_table)
+ # write c4 Lightning-Indexer keys BEFORE build_metadata so the scorer reads fresh+accumulated entries.
+ self.index_infer.write_indexer_k(x, infer_state, layer_weight, cos_table, sin_table)
+ return self.index_infer.build_metadata(x, q_lora, infer_state, layer_weight)
+
+ def _context_attention_kernel(
+ self,
+ q,
+ q_lora,
+ x,
+ infer_state: DeepseekV4InferStateInfo,
+ layer_weight: DeepseekV4TransformerLayerWeight,
+ out=None,
+ ):
+ meta = self._compress_and_index(q_lora, x, infer_state, layer_weight)
+ att_control = AttControl(
+ nsa_prefill=True,
+ nsa_prefill_dict={
+ "layer_index": self.layer_num_,
+ "compress_ratio": self.compress_ratio,
+ "head_dim_v": self.v_head_dim,
+ "softmax_scale": self.softmax_scale,
+ "attn_sink": layer_weight.attn_sink_.weight,
+ **meta,
+ },
+ )
+ attn_out = infer_state.prefill_att_state.prefill_att(
+ q=q,
+ k=infer_state.mem_manager.get_att_input_params(layer_index=self.layer_num_),
+ v=None,
+ att_control=att_control,
+ alloc_func=self.alloc_tensor,
+ out=out,
+ )
+ pad_q_len = getattr(infer_state, "_dsv4_prefill_pad_q_len", 0)
+ if pad_q_len:
+ # pad 行读 HOLD 槽位(参见 infer_struct._dsv4_prefill_pad_q_len),清零以保持确定性
+ attn_out[-pad_q_len:] = 0
+ return attn_out
+
+ # ------------------------------------------------------------------ attention (decode)
+ def token_attention_forward(
+ self, x, infer_state: DeepseekV4InferStateInfo, layer_weight: DeepseekV4TransformerLayerWeight
+ ):
+ q, q_lora, full_x = self._get_qkv(x, infer_state, layer_weight)
+ o = self._token_attention_kernel(q, q_lora, full_x, infer_state, layer_weight)
+ return self._get_o(o, infer_state, layer_weight)
+
+ def _token_attention_kernel(
+ self, q, q_lora, x, infer_state: DeepseekV4InferStateInfo, layer_weight: DeepseekV4TransformerLayerWeight
+ ):
+ self.compressor.prepare_states(x, infer_state, layer_weight)
+ self.compressor.fused_compress(infer_state, layer_weight, self.cos_compress_table, self.sin_compress_table)
+ self.index_infer.write_indexer_k(x, infer_state, layer_weight, self.cos_compress_table, self.sin_compress_table)
+ meta = self.index_infer.build_metadata(x, q_lora, infer_state, layer_weight)
+ att_control = AttControl(
+ nsa_decode=True,
+ nsa_decode_dict={
+ "layer_index": self.layer_num_,
+ "compress_ratio": self.compress_ratio,
+ "head_dim_v": self.v_head_dim,
+ "softmax_scale": self.softmax_scale,
+ "attn_sink": layer_weight.attn_sink_.weight,
+ **meta,
+ },
+ )
+ return infer_state.decode_att_state.decode_att(
+ q=q,
+ k=infer_state.mem_manager.get_att_input_params(layer_index=self.layer_num_),
+ v=None,
+ att_control=att_control,
+ )
+
+ # ------------------------------------------------------------------ moe
+ def _routed_experts(
+ self,
+ x,
+ weights,
+ indices,
+ infer_state: DeepseekV4InferStateInfo,
+ layer_weight: DeepseekV4TransformerLayerWeight,
+ ):
+ return layer_weight.experts_.experts_with_topk(
+ input_tensor=x,
+ topk_weights=weights,
+ topk_ids=indices,
+ is_prefill=infer_state.is_prefill,
+ clamp_limit=float(self.swiglu_limit),
+ alloc_tensor_func=self.alloc_tensor,
+ )
+
+ def _ffn_tp(self, input, infer_state: DeepseekV4InferStateInfo, layer_weight: DeepseekV4TransformerLayerWeight):
+ input = input.view(-1, self.embed_dim_)
+ gate_up = layer_weight.gate_up_proj.mm(input)
+ shared = self.alloc_tensor((input.size(0), gate_up.size(1) // 2), input.dtype)
+ silu_and_mul_fwd(gate_up, shared, limit=self.swiglu_limit)
+ input = None
+ gate_up = None
+ out = layer_weight.down_proj.mm(shared)
+ shared = None
+ return out
+
+ def _ffn(self, x, infer_state: DeepseekV4InferStateInfo, layer_weight: DeepseekV4TransformerLayerWeight):
+ x = x.view(-1, self.embed_dim_)
+ if not self.enable_ep_moe:
+ x = self._tpsp_allgather(input=x, infer_state=infer_state)
+
+ logits = layer_weight.gate_weight_.mm(x, out_dtype=torch.float32)
+ weights, indices = self._select_experts(logits, infer_state, layer_weight)
+ # shared expert 必须先于 routed 计算: fp8 路径 (FuseMoeTriton) 的 fused_experts
+ # 是 inplace 的，_routed_experts 返回后 x 已被覆盖为 routed 输出。
+ # DS4 shared experts also use the config swiglu_limit clamp, matching SGLang's
+ # DeepseekV2MLP(..., swiglu_limit=config.swiglu_limit) path.
+ shared = self._ffn_tp(input=x, infer_state=infer_state, layer_weight=layer_weight)
+ routed = self._routed_experts(x, weights, indices, infer_state, layer_weight)
+ if self.enable_ep_moe:
+ if self.tp_world_size_ > 1:
+ all_reduce(
+ shared,
+ op=dist.ReduceOp.SUM,
+ group=infer_state.dist_group,
+ async_op=False,
+ )
+ return routed + shared
+ out = routed + shared
+ return self._tpsp_reduce(input=out, infer_state=infer_state)
+
+ def _select_experts(
+ self, logits, infer_state: DeepseekV4InferStateInfo, layer_weight: DeepseekV4TransformerLayerWeight
+ ):
+ M = logits.shape[0]
+ bias = None
+ input_tokens = None
+ hash_indices_table = None
+ indices_dtype = torch.int64
+ if self.is_hash:
+ hash_indices_table = layer_weight.gate_tid2eid_.weight
+ if not hash_indices_table.is_contiguous():
+ hash_indices_table = hash_indices_table.contiguous()
+ indices_dtype = hash_indices_table.dtype
+ input_tokens = infer_state.input_ids.to(dtype=indices_dtype).contiguous()
+ else:
+ bias = layer_weight.gate_bias_.weight
+
+ weights = self.alloc_tensor((M, self.num_experts_per_tok), dtype=torch.float32, device=logits.device)
+ indices = self.alloc_tensor((M, self.num_experts_per_tok), dtype=indices_dtype, device=logits.device)
+ token_expert_indices = self.alloc_tensor((M, self.num_experts_per_tok), dtype=torch.int32, device=logits.device)
+ vllm_ops.topk_hash_softplus_sqrt(
+ weights,
+ indices,
+ token_expert_indices,
+ logits,
+ True,
+ self.routed_scaling_factor,
+ bias,
+ input_tokens,
+ hash_indices_table,
+ )
+ return weights, indices.long()
+
+
+class CompressorInfer:
+ """Window-softmax compressor. is_in_indexer=False compresses the c4/c128 latent KV into the
+ paged fp8 slab (attention extra_k); is_in_indexer=True reuses the SAME machinery (mirroring
+ sglang's Compressor(is_in_indexer=...)) with the indexer weights/dims/state pool to produce the
+ per-c4-entry Lightning-Indexer keys, emitted as dense bf16 (OUTPUT_BF16) then fp8-packed into
+ c4_indexer_pool by the caller. Indexer mode is c4-only."""
+
+ def __init__(self, layer_idx: int, network_config: dict, tp_world_size: int, is_in_indexer: bool = False):
+ super().__init__()
+ self.layer_idx_ = layer_idx
+ self.network_config_ = network_config
+ self.tp_world_size_ = tp_world_size
+ self.is_in_indexer = is_in_indexer
+ self.compress_ratio = network_config["compress_ratios"][layer_idx]
+ self.head_dim = network_config["head_dim"]
+ self.index_head_dim = network_config["index_head_dim"]
+ self.qk_rope_head_dim = network_config["qk_rope_head_dim"]
+ self.eps = network_config["rms_norm_eps"]
+ self._metadata = None
+
+ def prepare_states(
+ self,
+ x: torch.Tensor,
+ infer_state: DeepseekV4InferStateInfo,
+ layer_weight: DeepseekV4TransformerLayerWeight,
+ use_custom_tensor_manager: bool = True,
+ ):
+ # use_custom_tensor_manager=False routes the .mm outputs through torch.empty (stream-aware)
+ # instead of the stream-blind global cache -- required when this runs on the prefill aux stream.
+ self._metadata = prepare_compress_states(
+ infer_state=infer_state,
+ layer_idx=self.layer_idx_,
+ compress_ratio=self.compress_ratio,
+ is_in_indexer=self.is_in_indexer,
+ )
+ if self._metadata is not None:
+ if self.is_in_indexer:
+ # fused wkv/wgate GEMM -> [T, 2*coff*idx_hd] in the [kv | score] layout directly
+ # (same as the attention compressor_wkv_gate_).
+ self._metadata.kv_score = layer_weight.idx_cmp_wkv_gate_.mm(
+ x, use_custom_tensor_mananger=use_custom_tensor_manager, out_dtype=torch.float32
+ )
+ ape = layer_weight.idx_cmp_ape_.weight
+ else:
+ self._metadata.kv_score = layer_weight.compressor_wkv_gate_.mm(
+ x, use_custom_tensor_mananger=use_custom_tensor_manager, out_dtype=torch.float32
+ )
+ ape = layer_weight.compressor_ape_.weight
+ prepare_partial_states(
+ kv_score=self._metadata.kv_score,
+ metadata=self._metadata,
+ ape=ape,
+ compress_ratio=self.compress_ratio,
+ )
+ return self._metadata
+
+ def fused_compress(
+ self,
+ infer_state: DeepseekV4InferStateInfo,
+ layer_weight: DeepseekV4TransformerLayerWeight,
+ cos_table: torch.Tensor,
+ sin_table: torch.Tensor,
+ ):
+ if self.compress_ratio == 0:
+ return None
+ metadata = self._metadata
+ if metadata is None:
+ raise RuntimeError("DeepSeek-V4 compressor.prepare_states must run before fused_compress")
+ if self.is_in_indexer:
+ norm_weight = layer_weight.idx_cmp_norm_.weight
+ ape = layer_weight.idx_cmp_ape_.weight
+ head_dim = self.index_head_dim
+ else:
+ norm_weight = layer_weight.compressor_norm_.weight
+ ape = layer_weight.compressor_ape_.weight
+ head_dim = self.head_dim
+ return fused_compress_op(
+ kv_score=metadata.kv_score,
+ metadata=metadata,
+ norm_weight=norm_weight,
+ ape=ape,
+ eps=self.eps,
+ head_dim=head_dim,
+ qk_rope_head_dim=self.qk_rope_head_dim,
+ compress_ratio=self.compress_ratio,
+ cos_table=cos_table,
+ sin_table=sin_table,
+ output_bf16=self.is_in_indexer,
+ )
+
+
+class DeepseekV4IndexInfer:
+ """Model-side builder for the FlashMLA sparse-index metadata. Mirrors deepseek3_2's NsaInfer
+ boundary (the model owns ALL index construction; the attention backend only forwards final
+ tensors to flash_mla.flash_mla_with_kvcache) AND its c4 implementation: hadamard'd fp8 q/K, a
+ ragged gather of the compressed c4 keys, deep_gemm.fp8_mqa_logits, then topk -- adapted for the
+ replicated indexer (no gather-q/all_reduce), the c4-compressed entry space, and topk-512 (no
+ inheritance only because of those data-shape differences). swa metadata is precomputed in
+ init_some_extra_state; this class owns the c4 entry gather (build_compress_index) AND the c4
+ Lightning-Indexer scoring (gather + deep_gemm.fp8_mqa_logits + topk). Holds only static per-layer
+ config; all per-request data flows in via args. Invoke from _context/_token_attention_kernel
+ (after compressor.fused_compress, before *_att) so the c4 scorer/topk keep the same cuda-graph
+ capture position they had when this lived in the backend. The indexer is replicated (no TP collective)."""
+
+ def __init__(self, layer_idx: int, network_config: dict, tp_world_size: int):
+ self.layer_idx_ = layer_idx
+ self.compress_ratio = network_config["compress_ratios"][layer_idx]
+ self.index_topk = network_config["index_topk"]
+ self.index_head_dim = network_config["index_head_dim"]
+ self.qk_rope_head_dim = network_config["qk_rope_head_dim"]
+ self.index_n_heads = network_config["index_n_heads"]
+ self.tp_world_size_ = tp_world_size
+ self.indexer_score_scale = self.index_head_dim ** -0.5
+ self.indexer_weight_scale = self.indexer_score_scale * self.index_n_heads ** -0.5
+ # c4 layers own a second compressor (is_in_indexer) that writes the Lightning-Indexer key
+ # pool every step; _c4_indices gathers it back + scores via deep_gemm.fp8_mqa_logits.
+ self.indexer_compressor = (
+ CompressorInfer(layer_idx, network_config, tp_world_size, is_in_indexer=True)
+ if self.compress_ratio == 4
+ else None
+ )
+
+ def write_indexer_k(
+ self,
+ x,
+ infer_state: DeepseekV4InferStateInfo,
+ layer_weight,
+ cos_table,
+ sin_table,
+ use_custom_tensor_manager=True,
+ ):
+ """c4-only: compress this step's tokens into per-c4-entry indexer keys and pack them into
+ c4_indexer_pool. MUST run before build_metadata so the scorer (gather + deep_gemm.fp8_mqa_logits)
+ reads the finished entries; runs every step (incl. in the decode graph) so keys accumulate for
+ later long-context scoring. No-op on c128 / dense layers."""
+ if self.compress_ratio != 4:
+ return
+ self.indexer_compressor.prepare_states(
+ x, infer_state, layer_weight, use_custom_tensor_manager=use_custom_tensor_manager
+ )
+ self.indexer_compressor.fused_compress(infer_state, layer_weight, cos_table, sin_table)
+ scratch = self.indexer_compressor._metadata.out_buffer # [T, index_head_dim] bf16 (group-end rows valid)
+ # Rotate K (post norm+rope) by the SAME 1/sqrt(d) Hadamard the q kernel applies, so
+ # (Hq)·(Hk)=q·k (H orthogonal) and the fp8 quant of K stays accurate.
+ from lightllm.models.deepseek3_2.triton_kernel.hadamard_transform import hadamard_transform
+
+ scratch = hadamard_transform(scratch, scale=self.index_head_dim ** -0.5)
+ mem_manager = infer_state.mem_manager
+ positions = infer_state.position_ids
+ out_slots = mem_manager.full_to_c4_indexs[infer_state.mem_index.long().reshape(-1)]
+ # only group-end tokens finish a c4 entry; mask the rest to -1 so the packer skips them
+ # (mid-group tokens share the group's c4 slot -> avoids racing a finished slot).
+ completed = ((positions + 1) % 4 == 0) & (out_slots >= 0)
+ masked_slots = torch.where(completed, out_slots, torch.full_like(out_slots, -1)).to(torch.int32)
+ mem_manager.pack_indexer_k_to_cache(self.layer_idx_, masked_slots, scratch)
+
+ def build_metadata(
+ self, x, q_lora, infer_state: DeepseekV4InferStateInfo, layer_weight, use_custom_tensor_manager=True
+ ):
+ """Return the final flash_mla index tensors for this layer's compress variant. swa indices and
+ the per-token req_idx are layer-independent and precomputed once in init_some_extra_state
+ (read here); only the c4 scorer is per-layer. The backend pairs these with the
+ (data-independent, layer-keyed) fp8 cache-byte views it owns."""
+ swa_indices = infer_state.dsv4_swa_indices.unsqueeze(1)
+ swa_lengths = infer_state.dsv4_swa_lengths
+ positions = infer_state.position_ids
+ extra_indices = extra_lengths = None
+ if self.compress_ratio == 4:
+ idx_q_fp8, weights = self._indexer_q_weight(
+ x, q_lora, infer_state, layer_weight, use_custom_tensor_manager=use_custom_tensor_manager
+ )
+ extra_indices, extra_lengths = self._c4_indices(infer_state, idx_q_fp8, weights, positions)
+ elif self.compress_ratio == 128:
+ extra_indices = infer_state.dsv4_c128_indices.unsqueeze(1)
+ extra_lengths = infer_state.dsv4_c128_lengths
+ return {
+ "swa_indices": swa_indices,
+ "swa_lengths": swa_lengths,
+ "extra_indices": extra_indices,
+ "extra_lengths": extra_lengths,
+ }
+
+ def _indexer_q_weight(
+ self, x, q_lora, infer_state: DeepseekV4InferStateInfo, layer_weight, use_custom_tensor_manager=True
+ ):
+ """fp8 indexer q (mirrors deepseek3_2 NsaInfer): wq_b -> rope(last rope dims) -> 1/sqrt(d)
+ Hadamard -> per-token fp8 quant. Returns (idx_q_fp8 [T,H,d], weights [T,H]); the per-token q
+ fp8 scale and the head_dim^-0.5 * n_heads^-0.5 score scale are folded into weights -- the
+ deep_gemm.fp8_mqa_logits contract (fp8 q carries no companion scale). Replicated -> full heads."""
+ # Fused: wq_b mm -> rope(last rope dims) -> 1/sqrt(d) Hadamard -> per-token fp8 quant, with the
+ # per-token q scale + indexer_weight_scale folded into weights, all in ONE kernel (was 4 kernels:
+ # rotary_emb_fwd + hadamard_transform + act_quant + weights mul). freqs_cis is the compress rope
+ # table (same one the main compress-layer Q path uses); positions indexed inside the kernel.
+ from lightllm.models.deepseek_v4.triton_kernel.norm_rope_cuda import (
+ fused_q_indexer_rope_hadamard_quant,
+ )
+
+ token_num = q_lora.shape[0]
+ if x.shape[0] != token_num:
+ raise RuntimeError(
+ f"DeepSeek-V4 indexer expects full-token hidden states, got x={x.shape[0]} q_lora={token_num}"
+ )
+ idx_q = layer_weight.idx_wq_b_.mm(q_lora, use_custom_tensor_mananger=use_custom_tensor_manager).view(
+ token_num, self.index_n_heads, self.index_head_dim
+ )
+ raw_w = layer_weight.idx_weights_proj_.mm(x, use_custom_tensor_mananger=use_custom_tensor_manager).view(
+ token_num, self.index_n_heads
+ ) # [T, H] raw
+ idx_q_fp8, weights = fused_q_indexer_rope_hadamard_quant(
+ idx_q,
+ raw_w,
+ self.indexer_weight_scale,
+ self.freqs_cis,
+ infer_state.position_ids,
+ ) # fp8 [T,H,d]; weights [T,H,1] with q-scale + weight_scale folded
+ return idx_q_fp8, weights.squeeze(-1).contiguous()
+
+ def _c4_indices(self, infer_state: DeepseekV4InferStateInfo, idx_q_fp8, weights, positions):
+ """c4 scorer via the page-safe deep_gemm.fp8_paged_mqa_logits over the paged c4 indexer pool,
+ then masked topk-512 -> c4 slots. Fixed shapes (c4_cap pinned per graph bucket) keep the decode
+ cuda graph capturable."""
+ mem_manager = infer_state.mem_manager
+ workspace = infer_state.dsv4_workspace
+ index_topk = self.index_topk
+ max_entries = max(1, int(infer_state.max_kv_seq_len) // 4)
+ c4_cap = ((max_entries + 63) // 64) * 64
+
+ # entry space fits the budget -> every causal entry is selected; no scoring needed. The
+ # captured decode graph (graph_max_len -> max_entries > topk) always takes the scorer branch
+ # below, so this only shortcuts tiny eager contexts.
+ if max_entries <= index_topk:
+ from ..triton_kernel.build_compress_index_dsv4 import build_compress_index
+
+ slots, lengths = workspace.c4(infer_state.microbatch_index, positions.shape[0], c4_cap)
+ slots, lengths = build_compress_index(
+ infer_state.dsv4_sparse_req_idx,
+ positions,
+ infer_state.req_manager.req_to_token_indexs,
+ mem_manager.full_to_c4_indexs,
+ 4,
+ slots,
+ lengths,
+ )
+ return slots.unsqueeze(1), lengths
+
+ c4_len = torch.div(infer_state.b_seq_len, 4, rounding_mode="floor").to(torch.int32) # entries/req
+
+ device = positions.device
+ page_size = mem_manager.c4_indexer_pool.page_size
+
+ # The page table / row_page_table / valid_len / ctx_lens / paged-logits metadata / topk_lengths
+ # are LAYER-INDEPENDENT (depend on request layout + c4_cap, not on weights/layer). Build them on
+ # the first c4 layer of the forward and reuse on the other ~20 c4 layers (was rebuilt per layer:
+ # build_c4_indexer_page_table + a [T,npages] gather + clamp/reshape + get_paged_mqa_logits_metadata
+ # each, i.e. ~20x redundant index/copy/clamp launches). Lazy (not init_some_extra_state) so it is
+ # computed inside the decode cuda graph with the capture-forced shapes -> no graph-cap mismatch.
+ cached = getattr(infer_state, "_c4_paged_meta", None)
+ if cached is None:
+ from ..triton_kernel.gather_c4_indexer_k_dsv4 import build_c4_indexer_page_table
+
+ b_req_idx = infer_state.b_req_idx
+ batch = b_req_idx.shape[0]
+ page_table = build_c4_indexer_page_table(
+ mem_manager,
+ b_req_idx,
+ c4_len,
+ c4_cap,
+ infer_state.req_manager.req_to_token_indexs,
+ infer_state.req_manager.HOLD_REQUEST_ID,
+ )
+
+ if infer_state.is_prefill:
+ token_batch_pos = torch.repeat_interleave(
+ torch.arange(batch, device=device, dtype=torch.int32),
+ infer_state.b_q_seq_len,
+ output_size=positions.numel(),
+ )
+ row_page_table = page_table[token_batch_pos.long()].contiguous()
+ else:
+ row_page_table = page_table
+
+ valid_len = ((positions + 1) // 4).to(torch.int32)
+ ctx_lens = torch.clamp(valid_len, min=1).reshape(-1, 1).contiguous()
+ metadata = deep_gemm.get_paged_mqa_logits_metadata(
+ ctx_lens,
+ page_size,
+ deep_gemm.get_num_sms(),
+ )
+ topk_lengths = torch.clamp(
+ torch.minimum(valid_len, torch.full_like(valid_len, index_topk)), min=1
+ ).contiguous()
+ cached = (row_page_table, valid_len, ctx_lens, metadata, topk_lengths)
+ infer_state._c4_paged_meta = cached
+
+ row_page_table, valid_len, ctx_lens, metadata, topk_lengths = cached
+ kv_cache = mem_manager.c4_indexer_pool.get_layer_buffer(mem_manager.layer_to_c4_idx[self.layer_idx_]).view(
+ mem_manager.c4_indexer_pool.num_pages,
+ page_size,
+ 1,
+ self.index_head_dim + 4,
+ )
+ top_slots, _ = workspace.c4(infer_state.microbatch_index, idx_q_fp8.shape[0], index_topk)
+ if infer_state.is_prefill:
+ rows_per_chunk = max(1, _C4_PREFILL_LOGITS_BUDGET_BYTES // (c4_cap * 4))
+ if idx_q_fp8.shape[0] > rows_per_chunk:
+ for start in range(0, idx_q_fp8.shape[0], rows_per_chunk):
+ end = min(start + rows_per_chunk, idx_q_fp8.shape[0])
+ chunk_ctx_lens = ctx_lens[start:end]
+ self._c4_score_topk(
+ idx_q_fp8[start:end],
+ kv_cache,
+ weights[start:end],
+ chunk_ctx_lens,
+ row_page_table[start:end],
+ deep_gemm.get_paged_mqa_logits_metadata(
+ chunk_ctx_lens,
+ page_size,
+ deep_gemm.get_num_sms(),
+ ),
+ c4_cap,
+ valid_len[start:end],
+ top_slots[start:end],
+ page_size,
+ )
+ return top_slots.unsqueeze(1), topk_lengths
+
+ self._c4_score_topk(
+ idx_q_fp8,
+ kv_cache,
+ weights,
+ ctx_lens,
+ row_page_table,
+ metadata,
+ c4_cap,
+ valid_len,
+ top_slots,
+ page_size,
+ )
+ return top_slots.unsqueeze(1), topk_lengths
+
+ @staticmethod
+ def _c4_score_topk(
+ idx_q_fp8,
+ kv_cache,
+ weights,
+ ctx_lens,
+ row_page_table,
+ metadata,
+ c4_cap,
+ valid_len,
+ top_slots,
+ page_size,
+ ):
+ logits = deep_gemm.fp8_paged_mqa_logits(
+ idx_q_fp8.unsqueeze(1),
+ kv_cache,
+ weights,
+ ctx_lens,
+ row_page_table,
+ metadata,
+ c4_cap,
+ False,
+ )
+ topk_transform_512(
+ logits,
+ valid_len,
+ row_page_table,
+ top_slots,
+ page_size,
+ )
diff --git a/lightllm/models/deepseek_v4/layer_weights/__init__.py b/lightllm/models/deepseek_v4/layer_weights/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/lightllm/models/deepseek_v4/layer_weights/pre_and_post_layer_weight.py b/lightllm/models/deepseek_v4/layer_weights/pre_and_post_layer_weight.py
new file mode 100644
index 0000000000..54f29ce574
--- /dev/null
+++ b/lightllm/models/deepseek_v4/layer_weights/pre_and_post_layer_weight.py
@@ -0,0 +1,37 @@
+import torch
+from lightllm.common.basemodel import PreAndPostLayerWeight
+from lightllm.common.basemodel.layer_weights.meta_weights import (
+ EmbeddingWeight,
+ LMHeadWeight,
+ RMSNormWeight,
+ ParameterWeight,
+)
+
+
+class DeepseekV4PreAndPostLayerWeight(PreAndPostLayerWeight):
+ def __init__(self, data_type, network_config):
+ super().__init__(data_type, network_config)
+
+ hidden = network_config["hidden_size"]
+ vocab = network_config["vocab_size"]
+ hc_mult = network_config["hc_mult"]
+
+ # embeddings / lm_head / final norm (bf16, vocab tensor-parallel). V4 has no `model.` prefix
+ # and does not tie embeddings (tie_word_embeddings=false).
+ self.wte_weight_ = EmbeddingWeight(
+ dim=hidden, vocab_size=vocab, weight_name="embed.weight", data_type=self.data_type_
+ )
+ self.lm_head_weight_ = LMHeadWeight(
+ dim=hidden, vocab_size=vocab, weight_name="head.weight", data_type=self.data_type_
+ )
+ self.final_norm_weight_ = RMSNormWeight(dim=hidden, weight_name="norm.weight", data_type=self.data_type_)
+
+ # final hyper-connection head (collapses the hc_mult residual streams before the lm_head)
+ self.hc_head_fn_ = ParameterWeight(
+ weight_name="hc_head_fn", data_type=torch.float32, weight_shape=(hc_mult, hc_mult * hidden)
+ )
+ self.hc_head_base_ = ParameterWeight(
+ weight_name="hc_head_base", data_type=torch.float32, weight_shape=(hc_mult,)
+ )
+ self.hc_head_scale_ = ParameterWeight(weight_name="hc_head_scale", data_type=torch.float32, weight_shape=(1,))
+ return
diff --git a/lightllm/models/deepseek_v4/layer_weights/transformer_layer_weight.py b/lightllm/models/deepseek_v4/layer_weights/transformer_layer_weight.py
new file mode 100644
index 0000000000..072591a552
--- /dev/null
+++ b/lightllm/models/deepseek_v4/layer_weights/transformer_layer_weight.py
@@ -0,0 +1,339 @@
+import torch
+from lightllm.common.basemodel import TransformerLayerWeight
+from lightllm.common.basemodel.layer_weights.meta_weights import (
+ ROWMMWeight,
+ COLMMWeight,
+ ROWBMMWeight,
+ RMSNormWeight,
+ ParameterWeight,
+ TpAttSinkWeight,
+ FusedMoeWeight,
+)
+from ..triton_kernel.quant_convert import dequant_fp8_block_to_bf16
+
+
+class DeepseekV4TransformerLayerWeight(TransformerLayerWeight):
+ """Per-layer weights for DeepSeek-V4-Flash.
+
+ DS4 does not share DS2/DS3.2's ``model.layers.*.self_attn/mlp`` layout. Its attention is
+ HC + CSA, and routed experts are checkpointed as MXFP4 (fp4 release) or
+ FP8 block-128 (fp8 release, same layout as the dense fp8 weights).
+ """
+
+ def __init__(self, layer_num, data_type, network_config, quant_cfg=None):
+ super().__init__(layer_num, data_type, network_config, quant_cfg)
+ return
+
+ def _parse_config(self):
+ cfg = self.network_config_
+ self.hidden = cfg["hidden_size"]
+ self.n_heads = cfg["num_attention_heads"]
+ self.head_dim = cfg["head_dim"]
+ self.rope_dim = cfg["qk_rope_head_dim"]
+ self.q_lora_rank = cfg["q_lora_rank"]
+ self.o_lora_rank = cfg["o_lora_rank"]
+ self.o_groups = cfg["o_groups"]
+ self.index_n_heads = cfg["index_n_heads"]
+ self.index_head_dim = cfg["index_head_dim"]
+ self.n_routed_experts = cfg["n_routed_experts"]
+ self.moe_inter = cfg["moe_intermediate_size"]
+ self.num_hash_layers = cfg["num_hash_layers"]
+ self.vocab_size = cfg["vocab_size"]
+ self.hc_mult = cfg["hc_mult"]
+ self.mix_hc = (2 + self.hc_mult) * self.hc_mult
+ self.compress_ratio = cfg["compress_ratios"][self.layer_num_]
+ self.has_compressor = self.compress_ratio != 0
+ self.has_indexer = self.compress_ratio == 4
+ self.is_hash = self.layer_num_ < self.num_hash_layers
+ assert self.n_heads % self.tp_world_size_ == 0
+ assert self.o_groups % self.tp_world_size_ == 0
+ self.prefix = f"layers.{self.layer_num_}"
+
+ def _init_weight(self):
+ self._init_qkvo()
+ if self.has_compressor:
+ self._init_compressor()
+ if self.has_indexer:
+ self._init_indexer()
+ self._init_moe()
+ self._init_norm()
+ self._init_hyper_connection()
+
+ # ------------------------------------------------------------------ attention
+ def _init_qkvo(self):
+ p = f"{self.prefix}.attn"
+ # q low-rank A and kv (single replicated head) both consume the same attention input ->
+ # fuse into one fp8 GEMM; _get_qkv splits the [q_lora_rank | head_dim] output. (q_b is
+ # column-parallel over heads.)
+ self.wq_a_wkv_ = ROWMMWeight(
+ in_dim=self.hidden,
+ out_dims=[self.q_lora_rank, self.head_dim],
+ weight_names=[f"{p}.wq_a.weight", f"{p}.wkv.weight"],
+ data_type=self.data_type_,
+ quant_method=self.get_quant_method("wq_a"),
+ tp_rank=0,
+ tp_world_size=1,
+ )
+ self.wq_b_ = ROWMMWeight(
+ in_dim=self.q_lora_rank,
+ out_dims=[self.n_heads * self.head_dim],
+ weight_names=f"{p}.wq_b.weight",
+ data_type=self.data_type_,
+ quant_method=self.get_quant_method("wq_b"),
+ )
+ self.q_norm_ = RMSNormWeight(dim=self.q_lora_rank, weight_name=f"{p}.q_norm.weight", data_type=self.data_type_)
+ self.kv_norm_ = RMSNormWeight(dim=self.head_dim, weight_name=f"{p}.kv_norm.weight", data_type=self.data_type_)
+ self.attn_sink_ = TpAttSinkWeight(
+ all_q_head_num=self.n_heads, weight_name=f"{p}.attn_sink", data_type=torch.float32
+ )
+ # grouped low-rank output projection (wo_a per-group [in, o_lora], wo_b row-parallel
+ # [groups*o_lora -> hidden]).
+ per_group_in = self.n_heads * self.head_dim // self.o_groups
+ # When o_groups == tp_world_size (e.g. tp8) each rank owns exactly one group, so the
+ # grouped O-proj can collapse to a single GEMM. SGLang only enables the fp8 wo_a GEMM
+ # on Blackwell; on Hopper it keeps BF16, and this checkpoint stores wo_a as BF16.
+ # For >1 group per rank (tp < o_groups) the per-group inputs differ (block-diagonal), so
+ # keep the BF16 grouped bmm.
+ major, _ = torch.cuda.get_device_capability()
+ self.o_proj_fp8 = (self.o_groups // self.tp_world_size_) == 1 and major >= 10
+ if self.o_proj_fp8:
+ self.wo_a_ = ROWMMWeight(
+ in_dim=per_group_in,
+ out_dims=[self.o_groups * self.o_lora_rank],
+ weight_names=f"{p}.wo_a.weight",
+ data_type=self.data_type_,
+ quant_method=self.get_quant_method("wo_a"),
+ )
+ else:
+ self.wo_a_ = ROWBMMWeight(
+ dim0=self.o_groups,
+ dim1=per_group_in,
+ dim2=self.o_lora_rank,
+ weight_names=f"{p}.wo_a.weight",
+ data_type=self.data_type_,
+ quant_method=None,
+ )
+ self.wo_b_ = COLMMWeight(
+ in_dim=self.o_groups * self.o_lora_rank,
+ out_dims=[self.hidden],
+ weight_names=f"{p}.wo_b.weight",
+ data_type=self.data_type_,
+ quant_method=self.get_quant_method("wo_b"),
+ )
+
+ # ------------------------------------------------------------------ compressor / indexer
+ def _init_compressor(self):
+ prefix = f"{self.prefix}.attn.compressor"
+ head_dim = self.head_dim
+ ratio = self.compress_ratio
+
+ coff = 2 if ratio == 4 else 1
+ # wkv/wgate are bf16 (no scale) and replicated (single KV head).
+ self.compressor_wkv_gate_ = ROWMMWeight(
+ in_dim=self.hidden,
+ out_dims=[coff * head_dim, coff * head_dim],
+ weight_names=[f"{prefix}.wkv.weight", f"{prefix}.wgate.weight"],
+ data_type=self.data_type_,
+ quant_method=None,
+ tp_rank=0,
+ tp_world_size=1,
+ )
+ self.compressor_norm_ = RMSNormWeight(
+ dim=head_dim, weight_name=f"{prefix}.norm.weight", data_type=self.data_type_
+ )
+ self.compressor_ape_ = ParameterWeight(
+ weight_name=f"{prefix}.ape", data_type=torch.float32, weight_shape=(ratio, coff * head_dim)
+ )
+
+ def _init_indexer(self):
+ p = f"{self.prefix}.attn.indexer"
+ # The Lightning-Indexer is REPLICATED across TP ranks (like sglang/vllm), not head-sharded:
+ # q_lora and the attn input are already full on every rank, so each rank scores all
+ # index_n_heads locally and the c4 top-k is identical everywhere -- no gather/all_reduce.
+ # wq_b is FP8 in the checkpoint -> de-quantized to bf16 at load.
+ self.idx_wq_b_ = ROWMMWeight(
+ in_dim=self.q_lora_rank,
+ out_dims=[self.index_n_heads * self.index_head_dim],
+ weight_names=f"{p}.wq_b.weight",
+ data_type=self.data_type_,
+ quant_method=self.get_quant_method("idx_wq_b"),
+ tp_rank=0,
+ tp_world_size=1,
+ )
+ self.idx_weights_proj_ = ROWMMWeight(
+ in_dim=self.hidden,
+ out_dims=[self.index_n_heads],
+ weight_names=f"{p}.weights_proj.weight",
+ data_type=self.data_type_,
+ quant_method=None,
+ tp_rank=0,
+ tp_world_size=1,
+ )
+ coff = 2 # indexer compressor always uses ratio 4 (overlap)
+ # wkv/wgate share the same input -> one fused bf16 GEMM producing the [kv | gate] layout
+ # directly (same as the attention compressor_wkv_gate_).
+ self.idx_cmp_wkv_gate_ = ROWMMWeight(
+ in_dim=self.hidden,
+ out_dims=[coff * self.index_head_dim, coff * self.index_head_dim],
+ weight_names=[f"{p}.compressor.wkv.weight", f"{p}.compressor.wgate.weight"],
+ data_type=self.data_type_,
+ quant_method=None,
+ tp_rank=0,
+ tp_world_size=1,
+ )
+ self.idx_cmp_norm_ = RMSNormWeight(
+ dim=self.index_head_dim, weight_name=f"{p}.compressor.norm.weight", data_type=self.data_type_
+ )
+ self.idx_cmp_ape_ = ParameterWeight(
+ weight_name=f"{p}.compressor.ape", data_type=torch.float32, weight_shape=(4, coff * self.index_head_dim)
+ )
+
+ # ------------------------------------------------------------------ moe
+ def _init_moe(self):
+ p = f"{self.prefix}.ffn"
+ # Router gate weights stay bf16, but DS4 routing consumes fp32 GEMM output
+ # (SGLang linear_bf16_fp32 / vLLM router_logits_dtype=torch.float32).
+ self.gate_weight_ = ROWMMWeight(
+ in_dim=self.hidden,
+ out_dims=[self.n_routed_experts],
+ weight_names=f"{p}.gate.weight",
+ data_type=torch.bfloat16,
+ quant_method=None,
+ tp_rank=0,
+ tp_world_size=1,
+ )
+ if self.is_hash:
+ self.gate_tid2eid_ = ParameterWeight(
+ weight_name=f"{p}.gate.tid2eid",
+ data_type=torch.int64,
+ weight_shape=(self.vocab_size, self.network_config_["num_experts_per_tok"]),
+ )
+ else:
+ self.gate_bias_ = ParameterWeight(
+ weight_name=f"{p}.gate.bias", data_type=torch.float32, weight_shape=(self.n_routed_experts,)
+ )
+ # shared expert (dense, bf16 after de-quant): w1=gate, w3=up fused (row), w2=down (col).
+ # Named gate_up_proj/down_proj so the inherited Llama `_ffn_tp` (fused gate_up matmul +
+ # silu_and_mul triton kernel, no swiglu clamp) drives it directly. Order [w1, w3] = [gate, up]
+ # matches silu_and_mul_fwd's blocked layout (first half gate, second half up).
+ sp = f"{p}.shared_experts"
+ self.gate_up_proj = ROWMMWeight(
+ in_dim=self.hidden,
+ out_dims=[self.moe_inter, self.moe_inter],
+ weight_names=[f"{sp}.w1.weight", f"{sp}.w3.weight"],
+ data_type=self.data_type_,
+ quant_method=self.get_quant_method("shared_gate"),
+ )
+ self.down_proj = COLMMWeight(
+ in_dim=self.moe_inter,
+ out_dims=[self.hidden],
+ weight_names=f"{sp}.w2.weight",
+ data_type=self.data_type_,
+ quant_method=self.get_quant_method("shared_down"),
+ )
+ self.experts_ = FusedMoeWeight(
+ gate_proj_name="w1",
+ down_proj_name="w2",
+ up_proj_name="w3",
+ e_score_correction_bias_name="",
+ weight_prefix=f"{p}.experts",
+ n_routed_experts=self.n_routed_experts,
+ hidden_size=self.hidden,
+ moe_intermediate_size=self.moe_inter,
+ data_type=self.data_type_,
+ quant_method=self.quant_cfg.get_quant_method(self.layer_num_, "fused_moe"),
+ layer_num=self.layer_num_,
+ network_config=self.network_config_,
+ )
+
+ def _init_norm(self):
+ self.attn_norm_ = RMSNormWeight(
+ dim=self.hidden, weight_name=f"{self.prefix}.attn_norm.weight", data_type=self.data_type_
+ )
+ self.ffn_norm_ = RMSNormWeight(
+ dim=self.hidden, weight_name=f"{self.prefix}.ffn_norm.weight", data_type=self.data_type_
+ )
+
+ def _init_hyper_connection(self):
+ p = self.prefix
+ self.hc_attn_fn_ = ParameterWeight(
+ weight_name=f"{p}.hc_attn_fn",
+ data_type=torch.float32,
+ weight_shape=(self.mix_hc, self.hc_mult * self.hidden),
+ )
+ self.hc_attn_base_ = ParameterWeight(
+ weight_name=f"{p}.hc_attn_base", data_type=torch.float32, weight_shape=(self.mix_hc,)
+ )
+ self.hc_attn_scale_ = ParameterWeight(
+ weight_name=f"{p}.hc_attn_scale", data_type=torch.float32, weight_shape=(3,)
+ )
+ self.hc_ffn_fn_ = ParameterWeight(
+ weight_name=f"{p}.hc_ffn_fn",
+ data_type=torch.float32,
+ weight_shape=(self.mix_hc, self.hc_mult * self.hidden),
+ )
+ self.hc_ffn_base_ = ParameterWeight(
+ weight_name=f"{p}.hc_ffn_base", data_type=torch.float32, weight_shape=(self.mix_hc,)
+ )
+ self.hc_ffn_scale_ = ParameterWeight(
+ weight_name=f"{p}.hc_ffn_scale", data_type=torch.float32, weight_shape=(3,)
+ )
+
+ # ------------------------------------------------------------------ loading
+ def load_hf_weights(self, weights):
+ self._dequant_in_place(weights)
+ return super().load_hf_weights(weights)
+
+ def _fp8_scale_renames(self):
+ """Map weight name -> the scale name its quant method loads (e.g. `weight_scale_inv`
+ for DeepGEMM). Read from each MM weight's own `weight_scale_names`, so the rename
+ target always matches what that weight will look up; no-quant weights have None
+ entries and are skipped."""
+ renames = {}
+ for attr in self.__dict__.values():
+ weight_names = getattr(attr, "weight_names", ())
+ scale_names = getattr(attr, "weight_scale_names", ())
+ for weight_name, scale_name in zip(weight_names, scale_names):
+ if scale_name is not None:
+ renames[weight_name] = scale_name
+ return renames
+
+ def _dequant_in_place(self, weights):
+ p = self.prefix + "."
+ scale_renames = self._fp8_scale_renames()
+ # Convert every `.scale` belonging to this layer. Weights are loaded incrementally
+ # per safetensors shard, so the paired weight may live in another shard:
+ # - routed expert `.scale` follows the fused_moe quant method's weight_scale_suffix:
+ # MXFP4 consumes `.scale` as-is, FP8 DeepGEMM expects `.weight_scale_inv` (rename only);
+ # - FP8 matmul scales only need renaming for DeepGEMM, no weight required;
+ # - FP8 pairs on no-quant paths (wo_a's ROWBMMWeight) are expanded to bf16,
+ # the only case that truly requires weight and scale in the same shard.
+ expert_scale_suffix = self.experts_.quant_method.weight_scale_suffix
+ for scale_k in [k for k in list(weights.keys()) if k.startswith(p) and k.endswith(".scale")]:
+ if scale_k.startswith(f"{p}ffn.experts."):
+ if expert_scale_suffix is not None and expert_scale_suffix != "scale":
+ weights[scale_k[: -len("scale")] + expert_scale_suffix] = weights[scale_k].to(torch.float32)
+ del weights[scale_k]
+ continue
+ k = scale_k[: -len(".scale")] + ".weight"
+ target = scale_renames.get(k)
+ if target is not None: # FP8 e4m3, block-128 scale, run by DeepGEMM directly
+ weights[target] = weights[scale_k].to(torch.float32)
+ del weights[scale_k]
+ else:
+ weights[k] = dequant_fp8_block_to_bf16(weights[k], weights[scale_k]).to(self.data_type_)
+ del weights[scale_k]
+ # grouped-O (bf16 path only): reshape [groups*o_lora, in] -> [groups, in, o_lora] for the
+ # batched matmul. The fp8 path keeps wo_a as a plain [groups*o_lora, in] fp8 GEMM weight
+ # (its `.scale` is renamed to `.weight_scale_inv` by the loop above, not dequantized).
+ if not self.o_proj_fp8:
+ woa = f"{self.prefix}.attn.wo_a.weight"
+ if woa in weights and weights[woa].dim() == 2:
+ w = weights[woa]
+ per_group_in = self.n_heads * self.head_dim // self.o_groups
+ weights[woa] = w.view(self.o_groups, self.o_lora_rank, per_group_in).transpose(1, 2).contiguous()
+ # Keep c4 overlap APE in checkpoint layout [4, 2*head_dim]. SGLang reorders it
+ # because its compressor consumes ape.view(8, head_dim) by window offset. LightLLM
+ # adds APE into each token's two score halves before compression using position % 4,
+ # so the raw checkpoint layout is the equivalent representation here.
+ return
diff --git a/lightllm/models/deepseek_v4/model.py b/lightllm/models/deepseek_v4/model.py
new file mode 100644
index 0000000000..fd13bc43a0
--- /dev/null
+++ b/lightllm/models/deepseek_v4/model.py
@@ -0,0 +1,379 @@
+import copy
+import importlib.util
+import json
+import os
+
+import torch
+from lightllm.models.registry import ModelRegistry
+from lightllm.models.llama.model import LlamaTpPartModel
+from lightllm.common.basemodel.batch_objs import ModelInput
+from lightllm.common.req_manager import DeepseekV4ReqManager
+from lightllm.common.kv_cache_mem_manager import DeepseekV4MemoryManager
+from lightllm.models.deepseek_v4.layer_weights.pre_and_post_layer_weight import (
+ DeepseekV4PreAndPostLayerWeight,
+)
+from lightllm.models.deepseek_v4.layer_weights.transformer_layer_weight import (
+ DeepseekV4TransformerLayerWeight,
+)
+from lightllm.models.deepseek_v4.layer_infer.pre_layer_infer import (
+ DeepseekV4PreLayerInfer,
+)
+from lightllm.models.deepseek_v4.layer_infer.post_layer_infer import (
+ DeepseekV4PostLayerInfer,
+)
+from lightllm.models.deepseek_v4.layer_infer.transformer_layer_infer import (
+ DeepseekV4TransformerLayerInfer,
+)
+from lightllm.common.basemodel.attention import get_nsa_prefill_att_backend_class, get_nsa_decode_att_backend_class
+from lightllm.common.basemodel.attention.nsa.dsv4_fp8_flashmla_sparse import DSV4_NSA_BACKENDS
+from lightllm.models.deepseek_v4.infer_struct import DeepseekV4InferStateInfo
+from lightllm.models.deepseek_v4.workspace import DeepseekV4Workspace
+from lightllm.models.llama.yarn_rotary_utils import (
+ find_correction_range,
+ linear_ramp_mask,
+)
+from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num, get_env_start_args
+from lightllm.utils.log_utils import init_logger
+from lightllm.distributed.communication_op import dist_group_manager
+
+logger = init_logger(__name__)
+
+
+@ModelRegistry("deepseek_v4")
+class DeepseekV4TpPartModel(LlamaTpPartModel):
+ req_manager: DeepseekV4ReqManager
+ mem_manager: DeepseekV4MemoryManager
+
+ pre_and_post_weight_class = DeepseekV4PreAndPostLayerWeight
+ transformer_weight_class = DeepseekV4TransformerLayerWeight
+
+ pre_layer_infer_class = DeepseekV4PreLayerInfer
+ post_layer_infer_class = DeepseekV4PostLayerInfer
+ transformer_layer_infer_class = DeepseekV4TransformerLayerInfer
+
+ infer_state_class = DeepseekV4InferStateInfo
+
+ def _verify_params(self):
+ assert self.load_way == "HF", "only support HF format weights"
+ assert self.config["num_attention_heads"] % self.tp_world_size_ == 0
+ assert self.config["o_groups"] % self.tp_world_size_ == 0
+ assert self.config["index_n_heads"] % self.tp_world_size_ == 0
+ return
+
+ def _init_req_manager(self):
+ create_max_seq_len = 0
+ if self.batch_max_tokens is not None:
+ create_max_seq_len = max(create_max_seq_len, self.batch_max_tokens)
+ if self.max_seq_length is not None:
+ create_max_seq_len = max(create_max_seq_len, self.max_seq_length)
+
+ self._dsv4_req_manager_seq_len = create_max_seq_len
+ layer_num = self.config["n_layer"] + get_added_mtp_kv_layer_num()
+ self._dsv4_compress_rates = self._get_compress_rates(layer_num)
+ self.req_manager = DeepseekV4ReqManager(
+ self.max_req_num,
+ create_max_seq_len,
+ compress_rates=self._dsv4_compress_rates,
+ head_dim=self.config["head_dim"],
+ indexer_head_dim=self.config["index_head_dim"],
+ sliding_window=self.config["sliding_window"],
+ )
+ return
+
+ def _get_compress_rates(self, layer_num):
+ rates = list(self.config["compress_ratios"])
+ return rates[:layer_num]
+
+ def _init_mem_manager(self):
+ layer_num = self.config["n_layer"] + get_added_mtp_kv_layer_num()
+ compress_rates = getattr(self, "_dsv4_compress_rates", self._get_compress_rates(layer_num))
+ sliding_window = int(self.config["sliding_window"])
+ self.mem_manager = DeepseekV4MemoryManager(
+ self.max_total_token_num,
+ dtype=self.data_type,
+ head_num=1,
+ head_dim=self.config["head_dim"],
+ layer_num=layer_num,
+ compress_rates=compress_rates,
+ indexer_head_dim=self.config["index_head_dim"],
+ max_request_num=self.max_req_num,
+ sliding_window=sliding_window,
+ mem_fraction=self.mem_fraction,
+ )
+ self.req_manager.mem_manager = self.mem_manager
+ return
+
+ def _init_att_backend(self):
+ args = get_env_start_args()
+ if args.llm_kv_type == "None":
+ args.llm_kv_type = "fp8kv_dsa"
+ # TODO: 支持其他 kv type
+ if args.llm_kv_type != "fp8kv_dsa":
+ raise RuntimeError("DeepSeek-V4 requires llm_kv_type=fp8kv_dsa for packed FlashMLA sparse attention")
+ self.prefill_att_backend = get_nsa_prefill_att_backend_class(index=0, backend_map=DSV4_NSA_BACKENDS)(model=self)
+ self.decode_att_backend = get_nsa_decode_att_backend_class(index=0, backend_map=DSV4_NSA_BACKENDS)(model=self)
+
+ real_q_head_num = self.prefill_att_backend.real_q_head_num
+ padded_q_head_num = self.prefill_att_backend.padded_q_head_num
+ self.dsv4_workspace.init_flashmla_prefill_q(
+ real_q_head_num=real_q_head_num,
+ padded_q_head_num=padded_q_head_num,
+ head_dim=self.config["head_dim"],
+ dtype=self.data_type,
+ )
+ self.dsv4_workspace.init_flashmla_prefill_full_out(
+ q_head_num=padded_q_head_num,
+ head_dim_v=self.config["head_dim"],
+ dtype=self.data_type,
+ )
+ for layer_infer, layer_weight in zip(self.layers_infer, self.trans_layers_weight):
+ layer_infer.flashmla_q_head_num_ = padded_q_head_num
+ if padded_q_head_num == real_q_head_num:
+ continue
+ attn_sink = layer_weight.attn_sink_.weight
+ assert attn_sink.shape == (real_q_head_num,)
+ padded_attn_sink = torch.zeros((padded_q_head_num,), dtype=attn_sink.dtype, device=attn_sink.device)
+ padded_attn_sink[: attn_sink.shape[0]].copy_(attn_sink)
+ padded_attn_sink.load_ok = attn_sink.load_ok
+ layer_weight.attn_sink_.weight = padded_attn_sink
+ return
+
+ def _init_custom(self):
+ self._init_to_get_rotary()
+ self.dsv4_workspace = DeepseekV4Workspace(self)
+ if os.getenv("LIGHTLLM_DSV4_PREFILL_OVERLAP", "1") == "1":
+ prefill_aux_stream = torch.cuda.Stream()
+ for layer in self.layers_infer:
+ layer.dsv4_prefill_aux_stream = prefill_aux_stream
+ dist_group_manager.new_deepep_group(
+ self.config["n_routed_experts"],
+ self.config["hidden_size"],
+ self.config.get("num_experts_per_tok", 1),
+ self.config.get("moe_intermediate_size", self.config.get("intermediate_size")),
+ )
+ return
+
+ def _prepare_dsv4_slots(self, model_input: ModelInput) -> None:
+ """Commit DSV4 derived slots before BaseModel pads or scatters the generic input."""
+ if model_input.is_prefill and self.is_mtp_draft_model:
+ return
+ if model_input.mem_indexes_cpu is None:
+ return
+ if model_input.mem_indexes is None:
+ model_input.mem_indexes = model_input.mem_indexes_cpu.cuda(non_blocking=True)
+
+ if model_input.is_prefill:
+ self.req_manager.prepare_prefill(
+ b_req_idx_cpu=model_input.b_req_idx_cpu,
+ b_ready_cache_len_cpu=model_input.b_ready_cache_len,
+ b_seq_len_cpu=model_input.b_seq_len_cpu,
+ mem_indexes=model_input.mem_indexes,
+ )
+ return
+
+ if model_input.mtp_decode_slot_prepare_indices == ():
+ return
+ self.req_manager.prepare_decode(
+ model_input.b_req_idx_cpu,
+ model_input.b_seq_len_cpu,
+ model_input.b_mtp_index_cpu,
+ model_input.mem_indexes,
+ model_input.mtp_decode_slot_prepare_indices,
+ prepare_compress_slots=not self.is_mtp_draft_model,
+ )
+ return
+
+ @torch.no_grad()
+ def forward(self, model_input: ModelInput):
+ self._prepare_dsv4_slots(model_input)
+ return super().forward(model_input)
+
+ @torch.no_grad()
+ def microbatch_overlap_prefill(self, model_input0: ModelInput, model_input1: ModelInput):
+ self._prepare_dsv4_slots(model_input0)
+ self._prepare_dsv4_slots(model_input1)
+ return super().microbatch_overlap_prefill(model_input0, model_input1)
+
+ @torch.no_grad()
+ def microbatch_overlap_decode(self, model_input0: ModelInput, model_input1: ModelInput):
+ self._prepare_dsv4_slots(model_input0)
+ self._prepare_dsv4_slots(model_input1)
+ return super().microbatch_overlap_decode(model_input0, model_input1)
+
+ def _init_to_get_rotary(self):
+ # Interleaved (GPT-J) rope. Build complex64 freqs_cis tables (_freqs_cis_*) following the
+ # gemma4 two-variant convention; the fused CUDA Q/K kernels consume them directly, while
+ # _cos_cached_*/_sin_cached_* are .real/.imag views of the same storage for the inverse
+ # rope and compressor paths (deepseek2's interleaved triton rotary_emb_fwd).
+ # Sliding-window and compressed layers both use DeepSeek YaRN correction; only the
+ # RoPE base differs (rope_theta vs compress_rope_theta), matching SGLang/vLLM.
+ # Kept fp32 for accuracy (the apply upcasts anyway).
+ cfg = self.config
+ rs = cfg.get("rope_scaling", {}) or {}
+ dim = cfg["qk_rope_head_dim"]
+ # The rope tables MUST span every absolute position any request can produce (the served
+ # max_req_total_len / max_position_embeddings, up to 1M). Capping them shorter makes
+ # init_some_extra_state's index_select(cos/sin, position_ids) read OOB past the table at
+ # contexts beyond the cap (device-side assert / crash). ~268MB total at 1M, fp32x32 x4 views.
+ max_seq = max(int(self.max_seq_length), int(cfg.get("max_position_embeddings", 8192)))
+ freq_exponents = torch.arange(0, dim, 2, dtype=torch.float32, device="cuda") / dim
+ positions = torch.arange(max_seq, dtype=torch.float32, device="cuda")
+
+ rope_type = rs.get("rope_type", rs.get("type", "default"))
+ orig_max = rs.get("original_max_position_embeddings", 0)
+
+ def build_inv_freq(base):
+ freqs = 1.0 / (base ** freq_exponents)
+ if rope_type == "yarn" and orig_max > 0:
+ beta_fast = rs.get("beta_fast", 32)
+ beta_slow = rs.get("beta_slow", 1)
+ factor = rs.get("factor", 1)
+ if factor is None:
+ factor = cfg.get("max_position_embeddings", max_seq) / orig_max
+ low, high = find_correction_range(beta_fast, beta_slow, dim, base, orig_max)
+ smooth = 1 - linear_ramp_mask(low, high, dim // 2).cuda()
+ freqs = freqs / factor * (1 - smooth) + freqs * smooth
+ return freqs
+
+ sliding_freqs = build_inv_freq(cfg["rope_theta"])
+ f = torch.outer(positions, sliding_freqs) # [max_seq, dim//2]
+ self._freqs_cis_sliding = torch.complex(f.cos(), f.sin())
+
+ compress_freqs = build_inv_freq(cfg["compress_rope_theta"])
+ f = torch.outer(positions, compress_freqs) # [max_seq, dim//2]
+ self._freqs_cis_compress = torch.complex(f.cos(), f.sin())
+ self._cos_cached_sliding = self._freqs_cis_sliding.real
+ self._sin_cached_sliding = self._freqs_cis_sliding.imag
+ self._cos_cached_compress = self._freqs_cis_compress.real
+ self._sin_cached_compress = self._freqs_cis_compress.imag
+ # Each layer uses exactly one rope variant; wire its table once here (layers are already
+ # built: _init_infer_layer runs before _init_custom) instead of relaying via infer_state.
+ # The compressor needs the full compress tables (entry rope positions != token positions).
+ for layer in self.layers_infer:
+ layer.freqs_cis = self._freqs_cis_compress if layer.compress_ratio else self._freqs_cis_sliding
+ layer.cos_compress_table = self._cos_cached_compress
+ layer.sin_compress_table = self._sin_cached_compress
+ # the indexer-Q fused kernel (compress rope) needs the complex compress freqs table.
+ if getattr(layer, "index_infer", None) is not None:
+ layer.index_infer.freqs_cis = self._freqs_cis_compress
+ return
+
+
+class DeepSeekV4Tokenizer:
+ """Tokenizer wrapper for DeepSeek-V4's Python prompt encoding."""
+
+ # DeepSeek-V4 has a per-request thinking mode (...) toggled via
+ # chat_template_kwargs={"thinking": true}. It has no Jinja chat_template string,
+ # so advertise thinking support explicitly for tokenizer_supports_force_thinking().
+ supports_thinking = True
+
+ def __init__(self, tokenizer, model_dir):
+ self.tokenizer = tokenizer
+ self.model_dir = model_dir
+ self._encoding_module = None
+ self._added_vocab = None
+
+ def __getattr__(self, name):
+ return getattr(self.tokenizer, name)
+
+ def get_added_vocab(self):
+ if self._added_vocab is None:
+ self._added_vocab = self.tokenizer.get_added_vocab()
+ return self._added_vocab
+
+ def _get_encoding_module(self):
+ if self._encoding_module is not None:
+ return self._encoding_module
+
+ # Prefer the encoder shipped inside the model dir (respects any model-specific
+ # customization); fall back to the copy vendored in this repo, because some
+ # DeepSeek-V4 releases (e.g. the FP8 weights) do NOT ship an encoding/ dir.
+ # vLLM/sglang likewise vendor this encoder in-tree instead of depending on the
+ # model directory.
+ encoding_path = os.path.join(self.model_dir, "encoding", "encoding_dsv4.py")
+ if not os.path.exists(encoding_path):
+ encoding_path = os.path.join(os.path.dirname(__file__), "encoding", "encoding_dsv4.py")
+ if not os.path.exists(encoding_path):
+ raise FileNotFoundError(f"DeepSeek-V4 encoding file not found: {encoding_path}")
+
+ spec = importlib.util.spec_from_file_location("lightllm_deepseek_v4_encoding_dsv4", encoding_path)
+ if spec is None or spec.loader is None:
+ raise ImportError(f"failed to load DeepSeek-V4 encoding module from {encoding_path}")
+ module = importlib.util.module_from_spec(spec)
+ spec.loader.exec_module(module)
+ self._encoding_module = module
+ return module
+
+ def apply_chat_template(
+ self,
+ conversation=None,
+ messages=None,
+ tools=None,
+ tokenize=False,
+ add_generation_prompt=True,
+ thinking=None,
+ enable_thinking=None,
+ **kwargs,
+ ):
+ msgs = conversation if conversation is not None else messages
+ if msgs is None:
+ raise ValueError("Either 'conversation' or 'messages' must be provided")
+
+ msgs = copy.deepcopy(msgs)
+
+ # The model's DSML encoder (encode_arguments_to_dsml in encoding_dsv4.py) expects
+ # function.arguments as a JSON string and parses it internally. Upstream,
+ # build_prompt._normalize_tool_call_arguments converts arguments from the OpenAI
+ # JSON string to a dict (needed by Qwen3.x-style Jinja templates). A dict hits the
+ # encoder's except-branch and gets wrapped under a single name="arguments" param,
+ # which the model then imitates and amplifies across turns until required fields go
+ # missing. Re-serialize dicts back to JSON strings so the encoder emits one
+ # per real arg.
+ for msg in msgs:
+ content = msg.get("content")
+ if isinstance(content, list) and all(
+ isinstance(part, dict) and part.get("type") == "text" for part in content
+ ):
+ msg["content"] = "".join(part.get("text") or "" for part in content)
+ for tc in msg.get("tool_calls") or []:
+ fn = tc.get("function")
+ if isinstance(fn, dict) and isinstance(fn.get("arguments"), dict):
+ fn["arguments"] = json.dumps(fn["arguments"], ensure_ascii=False)
+
+ if tools:
+ wrapped_tools = []
+ for tool in tools:
+ if "function" in tool:
+ wrapped_tools.append(tool)
+ else:
+ wrapped_tools.append({"type": "function", "function": tool})
+
+ injected = False
+ for msg in msgs:
+ if msg.get("role") == "system":
+ existing = msg.get("tools") or []
+ msg["tools"] = existing + wrapped_tools
+ injected = True
+ break
+
+ if not injected:
+ msgs.insert(0, {"role": "system", "content": "", "tools": wrapped_tools})
+
+ if thinking is None:
+ thinking = bool(enable_thinking) if enable_thinking is not None else False
+ thinking_mode = "thinking" if thinking else "chat"
+ effort = kwargs.get("reasoning_effort")
+ if effort not in ("max", "high", None):
+ effort = None
+ encoding = self._get_encoding_module()
+ prompt = encoding.encode_messages(
+ msgs,
+ thinking_mode=thinking_mode,
+ drop_thinking=kwargs.get("drop_thinking", True),
+ add_default_bos_token=kwargs.get("add_default_bos_token", True),
+ reasoning_effort=effort,
+ )
+
+ if tokenize:
+ return self.tokenizer.encode(prompt, add_special_tokens=False)
+ return prompt
diff --git a/lightllm/models/deepseek_v4/triton_kernel/__init__.py b/lightllm/models/deepseek_v4/triton_kernel/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/lightllm/models/deepseek_v4/triton_kernel/build_compress_index_dsv4.py b/lightllm/models/deepseek_v4/triton_kernel/build_compress_index_dsv4.py
new file mode 100644
index 0000000000..dae5121ab9
--- /dev/null
+++ b/lightllm/models/deepseek_v4/triton_kernel/build_compress_index_dsv4.py
@@ -0,0 +1,79 @@
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _build_compress_index_kernel(
+ req_idx_ptr,
+ pos_ptr,
+ req_to_token_ptr,
+ req_to_token_stride0,
+ full_to_c_ptr,
+ index_ptr,
+ index_stride0,
+ length_ptr,
+ cap,
+ RATIO: tl.constexpr,
+ BLOCK_E: tl.constexpr,
+):
+ t = tl.program_id(0)
+ eb = tl.program_id(1)
+ req = tl.load(req_idx_ptr + t).to(tl.int64)
+ pos = tl.load(pos_ptr + t).to(tl.int64)
+ raw_len = (pos + 1) // RATIO
+
+ e = eb * BLOCK_E + tl.arange(0, BLOCK_E)
+ e_mask = e < cap
+ valid = (e < raw_len) & e_mask
+ # group-end token of compressed entry e: position e*RATIO + (RATIO-1).
+ end_pos = e * RATIO + (RATIO - 1)
+ safe_pos = tl.where(valid, end_pos, 0)
+ full_slot = tl.load(req_to_token_ptr + req * req_to_token_stride0 + safe_pos, mask=valid, other=0).to(tl.int64)
+ c_slot = tl.load(full_to_c_ptr + full_slot, mask=valid, other=-1).to(tl.int32)
+ tl.store(index_ptr + t * index_stride0 + e, c_slot, mask=e_mask)
+
+ if eb == 0:
+ tl.store(length_ptr + t, tl.maximum(raw_len, 1).to(tl.int32))
+
+
+def build_compress_index(
+ req_idx: torch.Tensor,
+ positions: torch.Tensor,
+ req_to_token_indexs: torch.Tensor,
+ full_to_c_indexs: torch.Tensor,
+ ratio: int,
+ index: torch.Tensor,
+ length: torch.Tensor,
+):
+ """Fused two-level group-end gather for the c4/c128 compressed-entry index tables.
+
+ For token t (at request `req_idx[t]`, absolute `positions[t]`) and compressed entry e:
+ slot[t, e] = full_to_c[ req_to_token[req, e*ratio + (ratio-1)] ] (the group-end token's full slot)
+ with slot = -1 where e >= (pos+1)//ratio (beyond the causal compressed length) or where the
+ full->c map is unset. Writes index [T, cap] and length [T] = clamp((pos+1)//ratio, 1).
+
+ Replaces the eager _gather_compress_slots/_c128/c4-causal torch chain. The caller owns the
+ output storage, so this wrapper does not allocate on the hot path.
+ """
+ T = positions.shape[0]
+ cap = index.shape[1]
+ if T == 0:
+ return index, length
+ BLOCK_E = 256
+ grid = (T, triton.cdiv(cap, BLOCK_E))
+ _build_compress_index_kernel[grid](
+ req_idx,
+ positions,
+ req_to_token_indexs,
+ req_to_token_indexs.stride(0),
+ full_to_c_indexs,
+ index,
+ index.stride(0),
+ length,
+ cap,
+ RATIO=ratio,
+ BLOCK_E=BLOCK_E,
+ num_warps=4,
+ )
+ return index, length
diff --git a/lightllm/models/deepseek_v4/triton_kernel/build_swa_index_dsv4.py b/lightllm/models/deepseek_v4/triton_kernel/build_swa_index_dsv4.py
new file mode 100644
index 0000000000..a1ef2d5be1
--- /dev/null
+++ b/lightllm/models/deepseek_v4/triton_kernel/build_swa_index_dsv4.py
@@ -0,0 +1,72 @@
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _build_swa_index_kernel(
+ req_idx_ptr,
+ pos_ptr,
+ req_to_token_ptr,
+ req_to_token_stride0,
+ full_to_swa_ptr,
+ swa_index_ptr,
+ swa_index_stride0,
+ swa_length_ptr,
+ WINDOW: tl.constexpr,
+ BLOCK_W: tl.constexpr,
+):
+ token_idx = tl.program_id(0)
+ req = tl.load(req_idx_ptr + token_idx).to(tl.int64)
+ pos = tl.load(pos_ptr + token_idx).to(tl.int64)
+
+ w = tl.arange(0, BLOCK_W)
+ w_mask = w < WINDOW
+ # most-recent-first window, identical to the eager _swa_indices (offset = position - arange).
+ offset = pos - w
+ valid = (offset >= 0) & w_mask
+ safe_offset = tl.where(valid, offset, 0)
+ full_slot = tl.load(req_to_token_ptr + req * req_to_token_stride0 + safe_offset, mask=valid, other=0).to(tl.int64)
+ swa_slot = tl.load(full_to_swa_ptr + full_slot, mask=valid, other=-1)
+ out = tl.where(valid, swa_slot, -1).to(tl.int32)
+ tl.store(swa_index_ptr + token_idx * swa_index_stride0 + w, out, mask=w_mask)
+
+ length = tl.minimum(tl.maximum(pos + 1, 1), WINDOW).to(tl.int32)
+ tl.store(swa_length_ptr + token_idx, length)
+
+
+def build_swa_index(
+ req_idx: torch.Tensor,
+ positions: torch.Tensor,
+ req_to_token_indexs: torch.Tensor,
+ full_to_swa_indexs: torch.Tensor,
+ swa_index: torch.Tensor,
+ swa_length: torch.Tensor,
+):
+ """Per-token sliding-window FlashMLA index table, built ONCE per forward (layer-independent:
+ full_to_swa is a single global map and the window is a model constant, so every layer's swa
+ indices are identical). Replaces DeepseekV4IndexInfer._swa_indices: for token t at
+ (req_idx, position) gather the last `window` tokens' full slots via req_to_token, then map
+ full -> swa; out-of-range positions store -1.
+
+ Writes (swa_index [T, window] int32, swa_length [T] int32). The caller owns the output storage;
+ the reader adds the s_q axis via unsqueeze(1).
+ """
+ T = positions.shape[0]
+ window = swa_index.shape[1]
+ if T == 0:
+ return swa_index, swa_length
+ _build_swa_index_kernel[(T,)](
+ req_idx,
+ positions,
+ req_to_token_indexs,
+ req_to_token_indexs.stride(0),
+ full_to_swa_indexs,
+ swa_index,
+ swa_index.stride(0),
+ swa_length,
+ WINDOW=window,
+ BLOCK_W=triton.next_power_of_2(window),
+ num_warps=4,
+ )
+ return swa_index, swa_length
diff --git a/lightllm/models/deepseek_v4/triton_kernel/csrc/norm_rope.cu b/lightllm/models/deepseek_v4/triton_kernel/csrc/norm_rope.cu
new file mode 100644
index 0000000000..b0a2963a53
--- /dev/null
+++ b/lightllm/models/deepseek_v4/triton_kernel/csrc/norm_rope.cu
@@ -0,0 +1,622 @@
+// Copyright 2023-2024 SGLang Team
+// SPDX-License-Identifier: Apache-2.0
+//
+// DeepSeek-V4 main Q/K and indexer-Q fused kernels.
+//
+// Adapted from SGLang commit 8cea0473ea5299bc04885f8f6ba71269415a39b5,
+// python/sglang/jit_kernel/csrc/deepseek_v4/main_norm_rope.cuh. This local
+// port keeps the device math, warp mapping, launch bounds, and Hopper PDL
+// protocol, while replacing tvm::ffi and the generic SGLang headers with a
+// small torch cpp_extension binding specialized for LightLLM's DSV4 shapes.
+
+#include
+#include
+#include
+#include
+
+#include
+#include
+#include
+
+#include
+#include
+#include
+
+namespace {
+
+constexpr uint32_t kWarpThreads = 32;
+constexpr uint32_t kFusedQBlockSize = 128;
+constexpr uint32_t kFusedQNumWarps = kFusedQBlockSize / kWarpThreads;
+constexpr uint32_t kFusedKBlockSize = 256;
+constexpr uint32_t kFusedKNumWarps = kFusedKBlockSize / kWarpThreads;
+
+constexpr int64_t kMainHeadDim = 512;
+constexpr int64_t kMainRopeDim = 64;
+constexpr int64_t kMainNopeDim = kMainHeadDim - kMainRopeDim;
+constexpr int64_t kIndexerHeadDim = 128;
+constexpr int64_t kIndexerRopeDim = 64;
+constexpr uint32_t kFlashMLAPageSize = 128;
+constexpr int32_t kFlashMLAPageBits = 7;
+constexpr int64_t kFlashMLADataBytes = 576;
+constexpr int64_t kFlashMLAScaleBytes = 8;
+constexpr int64_t kFlashMLABytesPerToken = kFlashMLADataBytes + kFlashMLAScaleBytes;
+constexpr int64_t kFlashMLAPageBytes =
+ ((kFlashMLABytesPerToken * kFlashMLAPageSize + kFlashMLADataBytes - 1) / kFlashMLADataBytes) *
+ kFlashMLADataBytes;
+
+template
+struct alignas(sizeof(T) * N) AlignedVector {
+ T data[N];
+
+ __device__ T& operator[](int i) { return data[i]; }
+ __device__ const T& operator[](int i) const { return data[i]; }
+};
+
+template
+__device__ __forceinline__ T warp_reduce_sum(T value) {
+#pragma unroll
+ for (int mask = 16; mask > 0; mask >>= 1) {
+ value += __shfl_xor_sync(0xffffffffu, value, mask, 32);
+ }
+ return value;
+}
+
+template
+__device__ __forceinline__ T warp_reduce_sum_width(T value) {
+ static_assert(NumThreads > 0 && NumThreads <= 32 && (NumThreads & (NumThreads - 1)) == 0);
+#pragma unroll
+ for (int mask = NumThreads / 2; mask > 0; mask >>= 1) {
+ value += __shfl_xor_sync(0xffffffffu, value, mask, 32);
+ }
+ return value;
+}
+
+template
+__device__ __forceinline__ T warp_reduce_max(T value) {
+#pragma unroll
+ for (int mask = 16; mask > 0; mask >>= 1) {
+ value = fmaxf(value, __shfl_xor_sync(0xffffffffu, value, mask, 32));
+ }
+ return value;
+}
+
+template
+__device__ __forceinline__ void pdl_wait_primary() {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+ if constexpr (UsePDL) {
+ asm volatile("griddepcontrol.wait;" ::: "memory");
+ }
+#endif
+}
+
+template
+__device__ __forceinline__ void pdl_trigger_secondary() {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+ if constexpr (UsePDL) {
+ asm volatile("griddepcontrol.launch_dependents;" :::);
+ }
+#endif
+}
+
+__device__ __forceinline__ int32_t cast_to_ue8m0(float x) {
+ const uint32_t u = __float_as_uint(x);
+ const int32_t exp = static_cast((u >> 23) & 0xffu);
+ const uint32_t mantissa = u & 0x7fffffu;
+ return exp + (mantissa != 0);
+}
+
+__device__ __forceinline__ float inv_scale_ue8m0(int32_t exp) {
+ return __uint_as_float(static_cast(127 + 127 - exp) << 23);
+}
+
+__device__ __forceinline__ uint16_t pack_fp8(float x, float y) {
+ constexpr float kFp8Max = 448.0f;
+ const float2 values = {
+ fmaxf(fminf(x, kFp8Max), -kFp8Max),
+ fmaxf(fminf(y, kFp8Max), -kFp8Max),
+ };
+ return __nv_cvt_float2_to_fp8x2(values, __NV_SATFINITE, __NV_E4M3);
+}
+
+template
+void launch_kernel(Kernel kernel, dim3 grid, dim3 block, cudaStream_t stream, bool enable_pdl, Args... args) {
+ cudaLaunchConfig_t config{};
+ config.gridDim = grid;
+ config.blockDim = block;
+ config.dynamicSmemBytes = 0;
+ config.stream = stream;
+
+ cudaLaunchAttribute attribute{};
+ if (enable_pdl) {
+ attribute.id = cudaLaunchAttributeProgrammaticStreamSerialization;
+ attribute.val.programmaticStreamSerializationAllowed = true;
+ config.attrs = &attribute;
+ config.numAttrs = 1;
+ }
+
+ C10_CUDA_CHECK(cudaLaunchKernelEx(&config, kernel, args...));
+}
+
+bool device_supports_pdl() {
+ return at::cuda::getCurrentDeviceProperties()->major >= 9;
+}
+
+void check_cuda_tensor(const at::Tensor& tensor, const char* name) {
+ TORCH_CHECK(tensor.is_cuda(), name, " must be a CUDA tensor");
+}
+
+void check_same_device(const at::Tensor& reference, const at::Tensor& tensor, const char* name) {
+ check_cuda_tensor(tensor, name);
+ TORCH_CHECK(tensor.get_device() == reference.get_device(), name, " must be on the same CUDA device as the input");
+}
+
+struct FusedQNormRopeParams {
+ const __nv_bfloat16* q_input;
+ __nv_bfloat16* q_output;
+ const float* freqs_cis;
+ const void* positions;
+ int64_t q_input_stride_batch;
+ int64_t q_output_stride_batch;
+ uint32_t batch_size;
+ uint32_t num_q_heads;
+ float eps;
+};
+
+template
+__global__ __launch_bounds__(kFusedQBlockSize, 16) void fused_q_norm_rope_kernel(
+ const __grid_constant__ FusedQNormRopeParams params) {
+ constexpr int64_t kVecSize = 8;
+ constexpr int64_t kLocalSize = kMainHeadDim / (kWarpThreads * kVecSize);
+ constexpr uint32_t kRopeVecs = kMainRopeDim / kVecSize;
+ using Storage = AlignedVector<__nv_bfloat16, kVecSize>;
+
+ const uint32_t warp_id = threadIdx.x / kWarpThreads;
+ const uint32_t lane_id = threadIdx.x % kWarpThreads;
+ const uint32_t work_id = blockIdx.x * kFusedQNumWarps + warp_id;
+ const uint32_t total_works = params.batch_size * params.num_q_heads;
+ if (work_id >= total_works) return;
+
+ const uint32_t batch_id = work_id / params.num_q_heads;
+ const uint32_t head_id = work_id % params.num_q_heads;
+ const auto* input_ptr =
+ params.q_input + batch_id * params.q_input_stride_batch + head_id * kMainHeadDim;
+ auto* output_ptr =
+ params.q_output + batch_id * params.q_output_stride_batch + head_id * kMainHeadDim;
+ const int32_t position = static_cast(static_cast(params.positions)[batch_id]);
+
+ __shared__ Storage rope_storage[kFusedQNumWarps][kRopeVecs];
+
+ pdl_wait_primary();
+
+ Storage input_vec[kLocalSize];
+#pragma unroll
+ for (int i = 0; i < kLocalSize; ++i) {
+ input_vec[i] = reinterpret_cast(input_ptr)[i * kWarpThreads + lane_id];
+ }
+
+ const float2 freq = reinterpret_cast(params.freqs_cis + position * kMainRopeDim)[lane_id];
+
+ float sum_of_squares = 0.0f;
+#pragma unroll
+ for (int i = 0; i < kLocalSize; ++i) {
+#pragma unroll
+ for (int j = 0; j < kVecSize; ++j) {
+ const float x = __bfloat162float(input_vec[i][j]);
+ sum_of_squares += x * x;
+ }
+ }
+ sum_of_squares = warp_reduce_sum(sum_of_squares);
+ const float norm_factor = rsqrtf(sum_of_squares / static_cast(kMainHeadDim) + params.eps);
+
+#pragma unroll
+ for (int i = 0; i < kLocalSize; ++i) {
+#pragma unroll
+ for (int j = 0; j < kVecSize; ++j) {
+ input_vec[i][j] = __float2bfloat16_rn(__bfloat162float(input_vec[i][j]) * norm_factor);
+ }
+ }
+
+ const bool is_rope_lane = lane_id >= kWarpThreads - kRopeVecs;
+#pragma unroll
+ for (int i = 0; i < kLocalSize; ++i) {
+ if (i == kLocalSize - 1 && is_rope_lane) {
+ rope_storage[warp_id][lane_id - (kWarpThreads - kRopeVecs)] = input_vec[i];
+ } else {
+ reinterpret_cast(output_ptr)[i * kWarpThreads + lane_id] = input_vec[i];
+ }
+ }
+ __syncwarp();
+
+ pdl_trigger_secondary();
+
+ const auto elem = reinterpret_cast(rope_storage[warp_id])[lane_id];
+ const float2 x = __bfloat1622float2(elem);
+ const float out_real = x.x * freq.x - x.y * freq.y;
+ const float out_imag = x.x * freq.y + x.y * freq.x;
+ reinterpret_cast<__nv_bfloat162*>(output_ptr + kMainNopeDim)[lane_id] =
+ __floats2bfloat162_rn(out_real, out_imag);
+}
+
+struct FusedKNormRopeFlashMLAParams {
+ const __nv_bfloat16* kv;
+ const __nv_bfloat16* kv_weight;
+ const float* freqs_cis;
+ const void* positions;
+ const int32_t* out_loc;
+ uint8_t* kvcache;
+ int64_t kv_stride_batch;
+ uint32_t batch_size;
+ float eps;
+};
+
+template
+__global__ __launch_bounds__(kFusedKBlockSize, 8) void fused_k_norm_rope_flashmla_kernel(
+ const __grid_constant__ FusedKNormRopeFlashMLAParams params) {
+ using Storage = AlignedVector<__nv_bfloat16, 2>;
+
+ const uint32_t tx = threadIdx.x;
+ const uint32_t warp_id = tx / kWarpThreads;
+ const uint32_t lane_id = tx % kWarpThreads;
+ const uint32_t work_id = blockIdx.x;
+ if (work_id >= params.batch_size) return;
+
+ const auto* input_ptr = params.kv + work_id * params.kv_stride_batch;
+ const int32_t position = static_cast(static_cast(params.positions)[work_id]);
+ const int32_t out_loc = params.out_loc[work_id];
+ const float* freqs_cis = params.freqs_cis + position * kMainRopeDim;
+
+ pdl_wait_primary();
+
+ const Storage input_vec = reinterpret_cast(input_ptr)[tx];
+ const Storage weight_vec = reinterpret_cast(params.kv_weight)[tx];
+ float2 data;
+ float2 freq{};
+ if (warp_id == kFusedKNumWarps - 1) {
+ freq = reinterpret_cast(freqs_cis)[lane_id];
+ }
+
+ float sum_of_squares = 0.0f;
+ const float input_x = __bfloat162float(input_vec[0]);
+ const float input_y = __bfloat162float(input_vec[1]);
+ sum_of_squares += input_x * input_x;
+ sum_of_squares += input_y * input_y;
+ const float warp_sum = warp_reduce_sum(sum_of_squares);
+
+ __shared__ float partial_sums[kFusedKNumWarps];
+ if (lane_id == 0) partial_sums[warp_id] = warp_sum;
+ __syncthreads();
+ sum_of_squares = warp_reduce_sum_width(partial_sums[lane_id % kFusedKNumWarps]);
+ const float norm_factor = rsqrtf(sum_of_squares / static_cast(kMainHeadDim) + params.eps);
+ data.x = input_x * norm_factor * __bfloat162float(weight_vec[0]);
+ data.y = input_y * norm_factor * __bfloat162float(weight_vec[1]);
+
+ const int32_t page = out_loc >> kFlashMLAPageBits;
+ const int32_t offset = out_loc & (kFlashMLAPageSize - 1);
+ auto* page_ptr = params.kvcache + static_cast(page) * kFlashMLAPageBytes;
+ auto* value_ptr = page_ptr + static_cast(offset) * kFlashMLADataBytes;
+
+ pdl_trigger_secondary();
+
+ if (warp_id == kFusedKNumWarps - 1) {
+ const float out_real = data.x * freq.x - data.y * freq.y;
+ const float out_imag = data.x * freq.y + data.y * freq.x;
+ reinterpret_cast<__nv_bfloat162*>(value_ptr + kMainNopeDim)[lane_id] =
+ __floats2bfloat162_rn(out_real, out_imag);
+ } else {
+ const float abs_max = warp_reduce_max(fmaxf(fabsf(data.x), fabsf(data.y)));
+ const float scale_raw = fmaxf(1.0e-4f, abs_max) / 448.0f;
+ const int32_t scale_ue8m0 = cast_to_ue8m0(scale_raw);
+ const float inv_scale = inv_scale_ue8m0(scale_ue8m0);
+ reinterpret_cast(value_ptr)[tx] = pack_fp8(data.x * inv_scale, data.y * inv_scale);
+ if (lane_id == 0) {
+ auto* scale_ptr = page_ptr + kFlashMLAPageSize * kFlashMLADataBytes + offset * kFlashMLAScaleBytes;
+ scale_ptr[warp_id] = static_cast(scale_ue8m0);
+ }
+ }
+}
+
+struct FusedQIndexerParams {
+ const __nv_bfloat16* q_input;
+ uint8_t* q_fp8;
+ const __nv_bfloat16* weight;
+ float* weights_out;
+ float weight_scale;
+ const float* freqs_cis;
+ const void* positions;
+ uint32_t batch_size;
+ uint32_t num_heads;
+};
+
+template
+__global__ __launch_bounds__(kFusedQBlockSize, 16) void fused_q_indexer_rope_hadamard_quant_kernel(
+ const __grid_constant__ FusedQIndexerParams params) {
+ constexpr int64_t kVecSize = 4;
+ constexpr uint32_t kRopeVecs = kIndexerRopeDim / kVecSize;
+ using Storage = AlignedVector<__nv_bfloat16, kVecSize>;
+ using Float4 = AlignedVector;
+
+ const uint32_t warp_id = threadIdx.x / kWarpThreads;
+ const uint32_t lane_id = threadIdx.x % kWarpThreads;
+ const uint32_t work_id = blockIdx.x * kFusedQNumWarps + warp_id;
+ const uint32_t total_works = params.batch_size * params.num_heads;
+ if (work_id >= total_works) return;
+
+ const uint32_t batch_id = work_id / params.num_heads;
+ const int32_t position = static_cast(static_cast(params.positions)[batch_id]);
+ const auto* input_ptr = params.q_input + static_cast(work_id) * kIndexerHeadDim;
+ const float* freqs_cis = params.freqs_cis + position * kIndexerRopeDim;
+ const bool is_rope_lane = lane_id >= kWarpThreads - kRopeVecs;
+
+ pdl_wait_primary();
+
+ const float weight_value = __bfloat162float(params.weight[work_id]);
+ const Storage input_vec = reinterpret_cast(input_ptr)[lane_id];
+ Float4 data;
+ Float4 freq{};
+#pragma unroll
+ for (int i = 0; i < kVecSize; ++i) data[i] = __bfloat162float(input_vec[i]);
+ if (is_rope_lane) {
+ freq = reinterpret_cast(freqs_cis)[lane_id - (kWarpThreads - kRopeVecs)];
+ const float x_real = data[0];
+ const float x_imag = data[1];
+ const float y_real = data[2];
+ const float y_imag = data[3];
+ data[0] = x_real * freq[0] - x_imag * freq[1];
+ data[1] = x_real * freq[1] + x_imag * freq[0];
+ data[2] = y_real * freq[2] - y_imag * freq[3];
+ data[3] = y_real * freq[3] + y_imag * freq[2];
+ }
+
+ pdl_trigger_secondary();
+
+ {
+ const float a0 = data[0], a1 = data[1], a2 = data[2], a3 = data[3];
+ data[0] = a0 + a1;
+ data[1] = a0 - a1;
+ data[2] = a2 + a3;
+ data[3] = a2 - a3;
+ }
+ {
+ const float a0 = data[0], a1 = data[1], a2 = data[2], a3 = data[3];
+ data[0] = a0 + a2;
+ data[1] = a1 + a3;
+ data[2] = a0 - a2;
+ data[3] = a1 - a3;
+ }
+#pragma unroll
+ for (uint32_t mask = 1; mask < kWarpThreads; mask <<= 1) {
+#pragma unroll
+ for (int i = 0; i < kVecSize; ++i) {
+ const float other = __shfl_xor_sync(0xffffffffu, data[i], mask, kWarpThreads);
+ data[i] = (lane_id & mask) ? (other - data[i]) : (data[i] + other);
+ }
+ }
+ constexpr float kHadamardScale = 0.08838834764831845f; // 1 / sqrt(128)
+#pragma unroll
+ for (int i = 0; i < kVecSize; ++i) data[i] *= kHadamardScale;
+
+ float local_max = fabsf(data[0]);
+#pragma unroll
+ for (int i = 1; i < kVecSize; ++i) local_max = fmaxf(local_max, fabsf(data[i]));
+ const float abs_max = warp_reduce_max(local_max);
+ const float scale = fmaxf(1.0e-4f, abs_max) / 448.0f;
+ const float inv_scale = 1.0f / scale;
+ AlignedVector result;
+ result[0] = pack_fp8(data[0] * inv_scale, data[1] * inv_scale);
+ result[1] = pack_fp8(data[2] * inv_scale, data[3] * inv_scale);
+ auto* output = reinterpret_cast*>(
+ params.q_fp8 + static_cast(work_id) * kIndexerHeadDim);
+ output[lane_id] = result;
+ params.weights_out[work_id] = weight_value * params.weight_scale * scale;
+}
+
+template
+void launch_q_norm(const FusedQNormRopeParams& params, cudaStream_t stream) {
+ const uint32_t works = params.batch_size * params.num_q_heads;
+ const dim3 grid((works + kFusedQNumWarps - 1) / kFusedQNumWarps);
+ launch_kernel(fused_q_norm_rope_kernel, grid, kFusedQBlockSize, stream, UsePDL, params);
+}
+
+template
+void launch_k_norm(const FusedKNormRopeFlashMLAParams& params, cudaStream_t stream) {
+ launch_kernel(
+ fused_k_norm_rope_flashmla_kernel, params.batch_size, kFusedKBlockSize, stream, UsePDL, params);
+}
+
+template
+void launch_indexer(const FusedQIndexerParams& params, cudaStream_t stream) {
+ const uint32_t works = params.batch_size * params.num_heads;
+ const dim3 grid((works + kFusedQNumWarps - 1) / kFusedQNumWarps);
+ launch_kernel(
+ fused_q_indexer_rope_hadamard_quant_kernel, grid, kFusedQBlockSize, stream, UsePDL, params);
+}
+
+void fused_q_norm_rope_cuda(
+ const at::Tensor& q_input,
+ const at::Tensor& q_output,
+ const at::Tensor& freqs_cis,
+ const at::Tensor& positions,
+ double eps) {
+ check_cuda_tensor(q_input, "q_input");
+ check_same_device(q_input, q_output, "q_output");
+ check_same_device(q_input, freqs_cis, "freqs_cis");
+ check_same_device(q_input, positions, "positions");
+ TORCH_CHECK(q_input.scalar_type() == at::kBFloat16, "q_input must be bfloat16");
+ TORCH_CHECK(q_output.scalar_type() == at::kBFloat16, "q_output must be bfloat16");
+ TORCH_CHECK(q_input.sizes() == q_output.sizes(), "q_input and q_output shapes must match");
+ TORCH_CHECK(q_input.dim() == 3 && q_input.size(2) == kMainHeadDim, "q_input must be [B, H, 512]");
+ TORCH_CHECK(q_input.stride(2) == 1 && q_input.stride(1) == kMainHeadDim, "q_input head rows must be contiguous");
+ TORCH_CHECK(
+ q_output.stride(2) == 1 && q_output.stride(1) == kMainHeadDim, "q_output head rows must be contiguous");
+ TORCH_CHECK(
+ freqs_cis.dim() == 2 && freqs_cis.size(1) == kMainRopeDim && freqs_cis.scalar_type() == at::kFloat &&
+ freqs_cis.is_contiguous(),
+ "freqs_cis must be contiguous [max_pos, 64] float32");
+ TORCH_CHECK(
+ positions.dim() == 1 && positions.size(0) == q_input.size(0) && positions.is_contiguous(),
+ "positions must be contiguous [B]");
+ TORCH_CHECK(
+ positions.scalar_type() == at::kInt || positions.scalar_type() == at::kLong,
+ "positions must be int32 or int64");
+ if (q_input.size(0) == 0) return;
+
+ c10::cuda::CUDAGuard guard(q_input.device());
+ const auto stream = at::cuda::getCurrentCUDAStream();
+ const FusedQNormRopeParams params{
+ reinterpret_cast(q_input.data_ptr()),
+ reinterpret_cast<__nv_bfloat16*>(q_output.data_ptr()),
+ freqs_cis.data_ptr(),
+ positions.data_ptr(),
+ q_input.stride(0),
+ q_output.stride(0),
+ static_cast(q_input.size(0)),
+ static_cast(q_input.size(1)),
+ static_cast(eps),
+ };
+ const bool use_pdl = device_supports_pdl();
+ if (positions.scalar_type() == at::kInt) {
+ use_pdl ? launch_q_norm(params, stream) : launch_q_norm(params, stream);
+ } else {
+ use_pdl ? launch_q_norm(params, stream) : launch_q_norm(params, stream);
+ }
+}
+
+void fused_k_norm_rope_flashmla_cuda(
+ const at::Tensor& kv,
+ const at::Tensor& kv_weight,
+ const at::Tensor& freqs_cis,
+ const at::Tensor& positions,
+ const at::Tensor& out_loc,
+ const at::Tensor& kvcache,
+ double eps,
+ int64_t page_size) {
+ check_cuda_tensor(kv, "kv");
+ check_same_device(kv, kv_weight, "kv_weight");
+ check_same_device(kv, freqs_cis, "freqs_cis");
+ check_same_device(kv, positions, "positions");
+ check_same_device(kv, out_loc, "out_loc");
+ check_same_device(kv, kvcache, "kvcache");
+ TORCH_CHECK(kv.dim() == 2 && kv.size(1) == kMainHeadDim && kv.stride(1) == 1, "kv must be [B, 512]");
+ TORCH_CHECK(kv.scalar_type() == at::kBFloat16, "kv must be bfloat16");
+ TORCH_CHECK(
+ kv_weight.dim() == 1 && kv_weight.size(0) == kMainHeadDim && kv_weight.scalar_type() == at::kBFloat16 &&
+ kv_weight.is_contiguous(),
+ "kv_weight must be contiguous [512] bfloat16");
+ TORCH_CHECK(
+ freqs_cis.dim() == 2 && freqs_cis.size(1) == kMainRopeDim && freqs_cis.scalar_type() == at::kFloat &&
+ freqs_cis.is_contiguous(),
+ "freqs_cis must be contiguous [max_pos, 64] float32");
+ TORCH_CHECK(
+ positions.dim() == 1 && positions.size(0) == kv.size(0) && positions.is_contiguous(),
+ "positions must be contiguous [B]");
+ TORCH_CHECK(
+ positions.scalar_type() == at::kInt || positions.scalar_type() == at::kLong,
+ "positions must be int32 or int64");
+ TORCH_CHECK(
+ out_loc.dim() == 1 && out_loc.size(0) == kv.size(0) && out_loc.scalar_type() == at::kInt &&
+ out_loc.is_contiguous(),
+ "out_loc must be contiguous [B] int32");
+ TORCH_CHECK(
+ kvcache.dim() == 2 && kvcache.scalar_type() == at::kByte && kvcache.is_contiguous() &&
+ kvcache.size(1) == kFlashMLAPageBytes,
+ "kvcache must be contiguous [num_pages, 74880] uint8");
+ TORCH_CHECK(page_size == kFlashMLAPageSize, "DSV4 main K CUDA kernel requires page_size=128");
+ if (kv.size(0) == 0) return;
+
+ c10::cuda::CUDAGuard guard(kv.device());
+ const auto stream = at::cuda::getCurrentCUDAStream();
+ const FusedKNormRopeFlashMLAParams params{
+ reinterpret_cast(kv.data_ptr()),
+ reinterpret_cast(kv_weight.data_ptr()),
+ freqs_cis.data_ptr(),
+ positions.data_ptr(),
+ out_loc.data_ptr(),
+ kvcache.data_ptr(),
+ kv.stride(0),
+ static_cast(kv.size(0)),
+ static_cast(eps),
+ };
+ const bool use_pdl = device_supports_pdl();
+ if (positions.scalar_type() == at::kInt) {
+ use_pdl ? launch_k_norm(params, stream) : launch_k_norm(params, stream);
+ } else {
+ use_pdl ? launch_k_norm(params, stream) : launch_k_norm(params, stream);
+ }
+}
+
+void fused_q_indexer_rope_hadamard_quant_cuda(
+ const at::Tensor& q_input,
+ const at::Tensor& q_fp8,
+ const at::Tensor& weight,
+ const at::Tensor& weights_out,
+ double weight_scale,
+ const at::Tensor& freqs_cis,
+ const at::Tensor& positions) {
+ check_cuda_tensor(q_input, "q_input");
+ check_same_device(q_input, q_fp8, "q_fp8");
+ check_same_device(q_input, weight, "weight");
+ check_same_device(q_input, weights_out, "weights_out");
+ check_same_device(q_input, freqs_cis, "freqs_cis");
+ check_same_device(q_input, positions, "positions");
+ TORCH_CHECK(
+ q_input.dim() == 3 && q_input.size(2) == kIndexerHeadDim && q_input.scalar_type() == at::kBFloat16 &&
+ q_input.is_contiguous(),
+ "q_input must be contiguous [B, H, 128] bfloat16");
+ TORCH_CHECK(
+ q_fp8.sizes() == q_input.sizes() && q_fp8.scalar_type() == at::kFloat8_e4m3fn && q_fp8.is_contiguous(),
+ "q_fp8 must be contiguous [B, H, 128] float8_e4m3fn");
+ TORCH_CHECK(
+ weight.dim() == 2 && weight.size(0) == q_input.size(0) && weight.size(1) == q_input.size(1) &&
+ weight.scalar_type() == at::kBFloat16 && weight.is_contiguous(),
+ "weight must be contiguous [B, H] bfloat16");
+ TORCH_CHECK(
+ weights_out.dim() == 3 && weights_out.size(0) == q_input.size(0) &&
+ weights_out.size(1) == q_input.size(1) && weights_out.size(2) == 1 &&
+ weights_out.scalar_type() == at::kFloat && weights_out.is_contiguous(),
+ "weights_out must be contiguous [B, H, 1] float32");
+ TORCH_CHECK(
+ freqs_cis.dim() == 2 && freqs_cis.size(1) == kIndexerRopeDim && freqs_cis.scalar_type() == at::kFloat &&
+ freqs_cis.is_contiguous(),
+ "freqs_cis must be contiguous [max_pos, 64] float32");
+ TORCH_CHECK(
+ positions.dim() == 1 && positions.size(0) == q_input.size(0) && positions.is_contiguous(),
+ "positions must be contiguous [B]");
+ TORCH_CHECK(
+ positions.scalar_type() == at::kInt || positions.scalar_type() == at::kLong,
+ "positions must be int32 or int64");
+ if (q_input.size(0) == 0) return;
+
+ c10::cuda::CUDAGuard guard(q_input.device());
+ const auto stream = at::cuda::getCurrentCUDAStream();
+ const FusedQIndexerParams params{
+ reinterpret_cast(q_input.data_ptr()),
+ reinterpret_cast(q_fp8.data_ptr()),
+ reinterpret_cast(weight.data_ptr()),
+ weights_out.data_ptr(),
+ static_cast(weight_scale),
+ freqs_cis.data_ptr(),
+ positions.data_ptr(),
+ static_cast(q_input.size(0)),
+ static_cast(q_input.size(1)),
+ };
+ const bool use_pdl = device_supports_pdl();
+ if (positions.scalar_type() == at::kInt) {
+ use_pdl ? launch_indexer(params, stream) : launch_indexer(params, stream);
+ } else {
+ use_pdl ? launch_indexer(params, stream) : launch_indexer(params, stream);
+ }
+}
+
+} // namespace
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) {
+ module.def("fused_q_norm_rope", &fused_q_norm_rope_cuda, "DSV4 fused main-Q RMSNorm + RoPE");
+ module.def(
+ "fused_k_norm_rope_flashmla",
+ &fused_k_norm_rope_flashmla_cuda,
+ "DSV4 fused main-K RMSNorm + RoPE + FlashMLA cache write");
+ module.def(
+ "fused_q_indexer_rope_hadamard_quant",
+ &fused_q_indexer_rope_hadamard_quant_cuda,
+ "DSV4 fused indexer-Q RoPE + Hadamard + FP8 quant");
+}
diff --git a/lightllm/models/deepseek_v4/triton_kernel/csrc/topk_transform.cu b/lightllm/models/deepseek_v4/triton_kernel/csrc/topk_transform.cu
new file mode 100644
index 0000000000..53d430243d
--- /dev/null
+++ b/lightllm/models/deepseek_v4/triton_kernel/csrc/topk_transform.cu
@@ -0,0 +1,345 @@
+// Copyright 2023-2024 SGLang Team
+// SPDX-License-Identifier: Apache-2.0
+//
+// DeepSeek-V4 c4-indexer top-k selection + page-translate.
+//
+// Adapted from SGLang commit 8cea0473ea5299bc04885f8f6ba71269415a39b5,
+// python/sglang/jit_kernel/csrc/deepseek_v4/topk_v1.cuh. LightLLM replaces
+// the tvm::ffi TensorView / TensorMatcher binding with a torch cpp_extension
+// launcher while preserving the original Hopper PDL protocol.
+
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+
+namespace {
+
+constexpr uint32_t kTopK = 512;
+constexpr uint32_t kTopKBlockSize = 512;
+constexpr uint32_t kSMEM = 16 * 1024 * sizeof(uint32_t); // 64KB (bytes)
+
+template
+__device__ __forceinline__ void pdl_wait_primary() {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+ if constexpr (UsePDL) asm volatile("griddepcontrol.wait;" ::: "memory");
+#endif
+}
+
+template
+__device__ __forceinline__ void pdl_trigger_secondary() {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+ if constexpr (UsePDL) asm volatile("griddepcontrol.launch_dependents;" :::);
+#endif
+}
+
+__device__ __forceinline__ uint8_t convert_to_uint8(float x) {
+ __half h = __float2half_rn(x);
+ uint16_t bits = __half_as_ushort(h);
+ uint16_t key = (bits & 0x8000) ? static_cast(~bits) : static_cast(bits | 0x8000);
+ return static_cast(key >> 8);
+}
+
+__device__ __forceinline__ uint32_t convert_to_uint32(float x) {
+ uint32_t bits = __float_as_uint(x);
+ return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u);
+}
+
+__device__ __forceinline__ int32_t page_to_indices(const int32_t* __restrict__ page_table, uint32_t i,
+ uint32_t page_bits) {
+ const uint32_t mask = (1u << page_bits) - 1u;
+ return (page_table[i >> page_bits] << page_bits) | (i & mask);
+}
+
+__device__ void naive_transform(const int32_t* __restrict__ page_table, int32_t* __restrict__ indices,
+ int32_t* __restrict__ raw_indices, const uint32_t length,
+ const uint32_t page_bits) {
+ if (const auto tx = threadIdx.x; tx < length) {
+ indices[tx] = page_to_indices(page_table, tx, page_bits);
+ if (raw_indices != nullptr) raw_indices[tx] = tx;
+ } else if (tx < kTopK) {
+ indices[tx] = -1; // fill invalid indices to -1
+ if (raw_indices != nullptr) raw_indices[tx] = -1;
+ }
+}
+
+__device__ void radix_topk(const float* __restrict__ input, int32_t* __restrict__ output,
+ const uint32_t length) {
+ constexpr uint32_t RADIX = 256;
+ constexpr uint32_t BLOCK_SIZE = kTopKBlockSize;
+ constexpr uint32_t SMEM_INPUT_SIZE = kSMEM / (2 * sizeof(int32_t));
+
+ alignas(128) __shared__ uint32_t _s_histogram_buf[2][RADIX + 32];
+ alignas(128) __shared__ uint32_t s_counter;
+ alignas(128) __shared__ uint32_t s_threshold_bin_id;
+ alignas(128) __shared__ uint32_t s_num_input[2];
+ alignas(128) __shared__ int32_t s_last_remain;
+
+ extern __shared__ uint32_t s_input_idx[][kSMEM / (2 * sizeof(int32_t))];
+
+ const uint32_t tx = threadIdx.x;
+ uint32_t remain_topk = kTopK;
+ auto& s_histogram = _s_histogram_buf[0];
+
+ const auto run_cumsum = [&] {
+#pragma unroll 8
+ for (int32_t i = 0; i < 8; ++i) {
+ static_assert(1 << 8 == RADIX);
+ if (tx < RADIX) {
+ const auto j = 1 << i;
+ const auto k = i & 1;
+ auto value = _s_histogram_buf[k][tx];
+ if (tx + j < RADIX) {
+ value += _s_histogram_buf[k][tx + j];
+ }
+ _s_histogram_buf[k ^ 1][tx] = value;
+ }
+ __syncthreads();
+ }
+ };
+
+ // stage 1: 8bit coarse histogram
+ if (tx < RADIX + 1) s_histogram[tx] = 0;
+ __syncthreads();
+ for (uint32_t idx = tx; idx < length; idx += BLOCK_SIZE) {
+ const auto bin = convert_to_uint8(input[idx]);
+ ::atomicAdd(&s_histogram[bin], 1);
+ }
+ __syncthreads();
+ run_cumsum();
+ if (tx < RADIX && s_histogram[tx] > remain_topk && s_histogram[tx + 1] <= remain_topk) {
+ s_threshold_bin_id = tx;
+ s_num_input[0] = 0;
+ s_counter = 0;
+ }
+ __syncthreads();
+
+ const auto threshold_bin = s_threshold_bin_id;
+ remain_topk -= s_histogram[threshold_bin + 1];
+ if (remain_topk == 0) {
+ for (uint32_t idx = tx; idx < length; idx += BLOCK_SIZE) {
+ const uint32_t bin = convert_to_uint8(input[idx]);
+ if (bin > threshold_bin) {
+ const auto pos = ::atomicAdd(&s_counter, 1);
+ output[pos] = idx;
+ }
+ }
+ __syncthreads();
+ return;
+ } else {
+ __syncthreads();
+ if (tx < RADIX + 1) {
+ s_histogram[tx] = 0;
+ }
+ __syncthreads();
+
+ for (uint32_t idx = tx; idx < length; idx += BLOCK_SIZE) {
+ const float raw_input = input[idx];
+ const uint32_t bin = convert_to_uint8(raw_input);
+ if (bin > threshold_bin) {
+ const auto pos = ::atomicAdd(&s_counter, 1);
+ output[pos] = idx;
+ } else if (bin == threshold_bin) {
+ const auto pos = ::atomicAdd(&s_num_input[0], 1);
+ if (pos < SMEM_INPUT_SIZE) {
+ s_input_idx[0][pos] = idx;
+ const auto sbin = convert_to_uint32(raw_input);
+ const auto sub_bin = (sbin >> 24) & 0xFF;
+ ::atomicAdd(&s_histogram[sub_bin], 1);
+ }
+ }
+ }
+ __syncthreads();
+ }
+
+ // stage 2: refine with 8bit radix passes
+#pragma unroll 4
+ for (int round = 0; round < 4; ++round) {
+ const auto r_idx = round % 2;
+
+ const auto raw_num_input = s_num_input[r_idx];
+ const auto num_input = raw_num_input < SMEM_INPUT_SIZE ? raw_num_input : SMEM_INPUT_SIZE;
+
+ run_cumsum();
+ if (tx < RADIX && s_histogram[tx] > remain_topk && s_histogram[tx + 1] <= remain_topk) {
+ s_threshold_bin_id = tx;
+ s_num_input[r_idx ^ 1] = 0;
+ s_last_remain = remain_topk - s_histogram[tx + 1];
+ }
+ __syncthreads();
+
+ const auto threshold_bin2 = s_threshold_bin_id;
+ remain_topk -= s_histogram[threshold_bin2 + 1];
+
+ if (remain_topk == 0) {
+ for (uint32_t i = tx; i < num_input; i += BLOCK_SIZE) {
+ const auto idx = s_input_idx[r_idx][i];
+ const auto offset = 24 - round * 8;
+ const auto bin = (convert_to_uint32(input[idx]) >> offset) & 0xFF;
+ if (bin > threshold_bin2) {
+ const auto pos = ::atomicAdd(&s_counter, 1);
+ output[pos] = idx;
+ }
+ }
+ __syncthreads();
+ break;
+ } else {
+ __syncthreads();
+ if (tx < RADIX + 1) {
+ s_histogram[tx] = 0;
+ }
+ __syncthreads();
+ for (uint32_t i = tx; i < num_input; i += BLOCK_SIZE) {
+ const auto idx = s_input_idx[r_idx][i];
+ const auto raw_input = input[idx];
+ const auto offset = 24 - round * 8;
+ const auto bin = (convert_to_uint32(raw_input) >> offset) & 0xFF;
+ if (bin > threshold_bin2) {
+ const auto pos = ::atomicAdd(&s_counter, 1);
+ output[pos] = idx;
+ } else if (bin == threshold_bin2) {
+ if (round == 3) {
+ const auto pos = ::atomicAdd(&s_last_remain, -1);
+ if (pos > 0) {
+ output[kTopK - pos] = idx;
+ }
+ } else {
+ const auto pos = ::atomicAdd(&s_num_input[r_idx ^ 1], 1);
+ if (pos < SMEM_INPUT_SIZE) {
+ s_input_idx[r_idx ^ 1][pos] = idx;
+ const auto sbin = convert_to_uint32(raw_input);
+ const auto sub_bin = (sbin >> (offset - 8)) & 0xFF;
+ ::atomicAdd(&s_histogram[sub_bin], 1);
+ }
+ }
+ }
+ }
+ __syncthreads();
+ }
+ }
+}
+
+struct TopKParams {
+ const float* scores;
+ const int32_t* seq_lens;
+ const int32_t* page_table;
+ int32_t* page_indices;
+ int32_t* raw_indices;
+ int64_t score_stride;
+ int64_t page_table_stride;
+ uint32_t page_bits;
+};
+
+template
+__global__ void topk_transform_kernel(const __grid_constant__ TopKParams params) {
+ const uint32_t work_id = blockIdx.x;
+ const uint32_t seq_len = params.seq_lens[work_id];
+ const auto score_ptr = params.scores + work_id * params.score_stride;
+ const auto page_ptr = params.page_table + work_id * params.page_table_stride;
+ const auto indices_ptr = params.page_indices + work_id * kTopK;
+ const auto raw_indices_ptr = params.raw_indices != nullptr ? params.raw_indices + work_id * kTopK : nullptr;
+ const uint32_t page_bits = params.page_bits;
+
+ pdl_wait_primary();
+
+ if (seq_len <= kTopK) {
+ naive_transform(page_ptr, indices_ptr, raw_indices_ptr, seq_len, page_bits);
+ } else {
+ __shared__ int32_t s_topk_indices[kTopK];
+ radix_topk(score_ptr, s_topk_indices, seq_len);
+ const auto tx = threadIdx.x;
+ indices_ptr[tx] = page_to_indices(page_ptr, s_topk_indices[tx], page_bits);
+ if (raw_indices_ptr != nullptr) raw_indices_ptr[tx] = s_topk_indices[tx];
+ }
+
+ pdl_trigger_secondary();
+}
+
+template
+void launch_topk(const TopKParams& params, uint32_t batch_size, cudaStream_t stream) {
+ constexpr uint32_t smem = kSMEM + sizeof(int32_t);
+ static const cudaError_t smem_result = cudaFuncSetAttribute(
+ topk_transform_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
+ C10_CUDA_CHECK(smem_result);
+
+ cudaLaunchConfig_t config{};
+ config.gridDim = batch_size;
+ config.blockDim = kTopKBlockSize;
+ config.dynamicSmemBytes = smem;
+ config.stream = stream;
+
+ cudaLaunchAttribute attribute{};
+ if constexpr (UsePDL) {
+ attribute.id = cudaLaunchAttributeProgrammaticStreamSerialization;
+ attribute.val.programmaticStreamSerializationAllowed = true;
+ config.attrs = &attribute;
+ config.numAttrs = 1;
+ }
+ C10_CUDA_CHECK(cudaLaunchKernelEx(&config, topk_transform_kernel, params));
+}
+
+} // namespace
+
+void topk_transform_512_cuda(at::Tensor scores, at::Tensor seq_lens, at::Tensor page_table,
+ at::Tensor page_indices, int64_t page_size, c10::optional raw_indices) {
+ TORCH_CHECK(scores.is_cuda(), "scores must be a CUDA tensor");
+ TORCH_CHECK(
+ seq_lens.is_cuda() && page_table.is_cuda() && page_indices.is_cuda() &&
+ seq_lens.get_device() == scores.get_device() && page_table.get_device() == scores.get_device() &&
+ page_indices.get_device() == scores.get_device(),
+ "all tensors must be on the same CUDA device");
+ TORCH_CHECK(scores.dim() == 2 && scores.dtype() == at::kFloat, "scores must be [B, S] float32");
+ TORCH_CHECK(
+ seq_lens.dim() == 1 && seq_lens.size(0) == scores.size(0) && seq_lens.dtype() == at::kInt &&
+ seq_lens.is_contiguous(),
+ "seq_lens must be [B] int32 contiguous");
+ TORCH_CHECK(
+ page_table.dim() == 2 && page_table.size(0) == scores.size(0) && page_table.dtype() == at::kInt,
+ "page_table must be [B, P] int32");
+ TORCH_CHECK(page_indices.dim() == 2 && page_indices.dtype() == at::kInt && page_indices.is_contiguous(),
+ "page_indices must be [B, 512] int32 contiguous");
+ TORCH_CHECK(page_indices.size(0) == scores.size(0), "page_indices first dim must match scores");
+ TORCH_CHECK(page_indices.size(1) == (int64_t)kTopK, "page_indices second dim must be 512");
+ TORCH_CHECK(scores.stride(1) == 1 && page_table.stride(1) == 1, "scores/page_table last dim must be contiguous");
+ TORCH_CHECK(page_size > 0 && (page_size & (page_size - 1)) == 0, "page_size must be a power of two");
+
+ const uint32_t batch_size = scores.size(0);
+ if (batch_size == 0) return;
+ const uint32_t page_bits = __builtin_ctzll(page_size);
+
+ int32_t* raw_ptr = nullptr;
+ if (raw_indices.has_value()) {
+ auto& r = raw_indices.value();
+ TORCH_CHECK(
+ r.is_cuda() && r.get_device() == scores.get_device() && r.dim() == 2 && r.size(0) == scores.size(0) &&
+ r.size(1) == (int64_t)kTopK && r.dtype() == at::kInt && r.is_contiguous(),
+ "raw_indices must be [B, 512] int32 contiguous on the same CUDA device");
+ raw_ptr = r.data_ptr();
+ }
+
+ c10::cuda::CUDAGuard guard(scores.device());
+ auto stream = at::cuda::getCurrentCUDAStream();
+ const TopKParams params{
+ scores.data_ptr(),
+ seq_lens.data_ptr(),
+ page_table.data_ptr(),
+ page_indices.data_ptr(),
+ raw_ptr,
+ scores.stride(0),
+ page_table.stride(0),
+ page_bits,
+ };
+ if (at::cuda::getCurrentDeviceProperties()->major >= 9) {
+ launch_topk(params, batch_size, stream);
+ } else {
+ launch_topk(params, batch_size, stream);
+ }
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+ m.def("topk_transform_512", &topk_transform_512_cuda, "DeepSeek-V4 c4 indexer top-512 + page translate");
+}
diff --git a/lightllm/models/deepseek_v4/triton_kernel/destindex_copy_indexer_k_dsv4.py b/lightllm/models/deepseek_v4/triton_kernel/destindex_copy_indexer_k_dsv4.py
new file mode 100644
index 0000000000..3510b92c30
--- /dev/null
+++ b/lightllm/models/deepseek_v4/triton_kernel/destindex_copy_indexer_k_dsv4.py
@@ -0,0 +1,92 @@
+import torch
+
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _fwd_kernel_destindex_copy_indexer_k_dsv4(
+ K,
+ Dest_loc,
+ O_fp8,
+ O_f32,
+ stride_k_bs,
+ stride_k_d,
+ FP8_MIN: tl.constexpr,
+ FP8_MAX: tl.constexpr,
+ SCALE_MIN: tl.constexpr,
+ HEAD_DIM: tl.constexpr,
+ PAGE_SIZE: tl.constexpr,
+ BYTES_PER_PAGE: tl.constexpr,
+):
+ cur_index = tl.program_id(0)
+ dest_index = tl.load(Dest_loc + cur_index).to(tl.int64)
+ # negative dest (unmapped slot) is a no-op, not an OOB write into a neighboring page.
+ if dest_index < 0:
+ return
+
+ page = dest_index // PAGE_SIZE
+ token_in_page = dest_index % PAGE_SIZE
+
+ offs_d = tl.arange(0, HEAD_DIM)
+ vals = tl.load(K + cur_index * stride_k_bs + offs_d * stride_k_d).to(tl.float32)
+ amax = tl.max(tl.abs(vals), axis=0)
+ # per-token plain fp32 scale (not ue8m0), matching DeepseekV4MemoryManager._pack_indexer_k
+ scale = tl.maximum(amax / FP8_MAX, SCALE_MIN)
+ k_fp8 = tl.clamp(vals / scale, min=FP8_MIN, max=FP8_MAX).to(tl.float8e4nv)
+
+ data_base = page * BYTES_PER_PAGE + token_in_page * HEAD_DIM
+ tl.store(O_fp8 + data_base + offs_d, k_fp8)
+ scale_idx = (page * BYTES_PER_PAGE + PAGE_SIZE * HEAD_DIM) // 4 + token_in_page
+ tl.store(O_f32 + scale_idx, scale)
+ return
+
+
+@torch.no_grad()
+def destindex_copy_indexer_k_dsv4(
+ K: torch.Tensor,
+ DestLoc: torch.Tensor,
+ O_buffer: torch.Tensor,
+ page_size: int,
+):
+ """Packed indexer-K page-slab writer (DeepSeek-V4 c4/CSA layers).
+
+ K: [T, 128] bf16 unquantized indexer keys.
+ DestLoc: [T] int — c4-pool-local token slots; must already be allocated by the caller.
+ Negative slots (unmapped) are skipped.
+ O_buffer: [num_pages, bytes_per_page] uint8 — one layer's slab from the c4 indexer
+ PackedPagePool (128B fp8 data region + 4B fp32 scale tail per token).
+
+ Bit-compatible with DeepseekV4MemoryManager._pack_indexer_k + PackedPagePool.write.
+ """
+ seq_len = DestLoc.shape[0]
+ if seq_len == 0:
+ return
+ head_dim, scale_bytes = 128, 4
+
+ K = K.reshape(-1, head_dim)
+ assert K.shape[0] == seq_len, f"Expected K shape[0]={seq_len}, got {K.shape[0]}"
+ assert K.dtype == torch.bfloat16, f"Expected bf16 indexer K, got {K.dtype}"
+ bytes_per_page = O_buffer.shape[-1]
+ assert O_buffer.dtype == torch.uint8 and O_buffer.is_contiguous()
+ assert bytes_per_page % 4 == 0
+ assert bytes_per_page >= page_size * (head_dim + scale_bytes)
+
+ flat = O_buffer.view(-1)
+ _fwd_kernel_destindex_copy_indexer_k_dsv4[(seq_len,)](
+ K,
+ DestLoc,
+ flat.view(torch.float8_e4m3fn),
+ flat.view(torch.float32),
+ K.stride(0),
+ K.stride(1),
+ FP8_MIN=torch.finfo(torch.float8_e4m3fn).min,
+ FP8_MAX=torch.finfo(torch.float8_e4m3fn).max,
+ SCALE_MIN=1e-4,
+ HEAD_DIM=head_dim,
+ PAGE_SIZE=page_size,
+ BYTES_PER_PAGE=bytes_per_page,
+ num_warps=1,
+ num_stages=1,
+ )
+ return
diff --git a/lightllm/models/deepseek_v4/triton_kernel/destindex_copy_kv_flashmla_dsv4.py b/lightllm/models/deepseek_v4/triton_kernel/destindex_copy_kv_flashmla_dsv4.py
new file mode 100644
index 0000000000..a3ec6ed8cf
--- /dev/null
+++ b/lightllm/models/deepseek_v4/triton_kernel/destindex_copy_kv_flashmla_dsv4.py
@@ -0,0 +1,121 @@
+import torch
+
+import triton
+import triton.language as tl
+from triton.language.extra import libdevice
+
+
+@triton.jit
+def _fwd_kernel_destindex_copy_kv_flashmla_dsv4(
+ KV,
+ Dest_loc,
+ O_fp8,
+ O_bf16,
+ O_u8,
+ stride_kv_bs,
+ stride_kv_d,
+ FP8_MIN: tl.constexpr,
+ FP8_MAX: tl.constexpr,
+ SCALE_MIN: tl.constexpr,
+ NOPE_DIM: tl.constexpr,
+ ROPE_DIM: tl.constexpr,
+ GROUP_SIZE: tl.constexpr,
+ NUM_GROUPS: tl.constexpr,
+ SCALE_BYTES: tl.constexpr,
+ PAGE_SIZE: tl.constexpr,
+ BYTES_PER_PAGE: tl.constexpr,
+):
+ cur_index = tl.program_id(0)
+ dest_index = tl.load(Dest_loc + cur_index).to(tl.int64)
+ # negative dest (unmapped slot, e.g. full_to_c* rows that never closed a group) is a no-op,
+ # not an OOB write into a neighboring page.
+ if dest_index < 0:
+ return
+
+ page = dest_index // PAGE_SIZE
+ token_in_page = dest_index % PAGE_SIZE
+ data_base = page * BYTES_PER_PAGE + token_in_page * (NOPE_DIM + ROPE_DIM * 2)
+ scale_base = page * BYTES_PER_PAGE + PAGE_SIZE * (NOPE_DIM + ROPE_DIM * 2) + token_in_page * SCALE_BYTES
+
+ # nope: per-group ue8m0 quant. SCALE_BYTES(=NUM_GROUPS+1) lanes cover the exponent bytes
+ # plus the trailing zero pad byte in one store. libdevice.log2 (not tl.log2, which is the
+ # approx instruction) and the bit-packed 2**e keep this bit-exact with the torch oracle
+ # DeepseekV4MemoryManager._pack_mla_kv.
+ offs_g = tl.arange(0, SCALE_BYTES)
+ offs_e = tl.arange(0, GROUP_SIZE)
+ group_mask = offs_g < NUM_GROUPS
+ kv_ptrs = KV + cur_index * stride_kv_bs + (offs_g[:, None] * GROUP_SIZE + offs_e[None, :]) * stride_kv_d
+ vals = tl.load(kv_ptrs, mask=group_mask[:, None], other=0.0).to(tl.float32)
+ amax = tl.max(tl.abs(vals), axis=1)
+ scale_exp = tl.ceil(libdevice.log2(tl.maximum(amax / FP8_MAX, SCALE_MIN))).to(tl.int32)
+ scale = ((scale_exp + 127) << 23).to(tl.float32, bitcast=True)
+ kv_fp8 = tl.clamp(vals / scale[:, None], min=FP8_MIN, max=FP8_MAX).to(tl.float8e4nv)
+ tl.store(O_fp8 + data_base + offs_g[:, None] * GROUP_SIZE + offs_e[None, :], kv_fp8, mask=group_mask[:, None])
+ scale_bytes = tl.where(group_mask, scale_exp + 127, 0).to(tl.uint8)
+ tl.store(O_u8 + scale_base + offs_g, scale_bytes)
+
+ # rope: bf16 passthrough into the data region right after the nope bytes
+ offs_r = tl.arange(0, ROPE_DIM)
+ rope = tl.load(KV + cur_index * stride_kv_bs + (NOPE_DIM + offs_r) * stride_kv_d)
+ tl.store(O_bf16 + (data_base + NOPE_DIM) // 2 + offs_r, rope)
+ return
+
+
+@torch.no_grad()
+def destindex_copy_kv_flashmla_dsv4(
+ KV: torch.Tensor,
+ DestLoc: torch.Tensor,
+ O_buffer: torch.Tensor,
+ page_size: int,
+):
+ """fp8_ds_mla packed page-slab writer (DeepSeek-V4 ABI, all latent pools).
+
+ KV: [T, 512] bf16 — 448 normed-latent dims + 64 rope'd dims per token.
+ DestLoc: [T] int — pool-local token slots (page = slot // page_size); the pool HOLD slot is
+ a valid in-bounds row, negative slots (unmapped) are skipped. Slots must already be
+ resolved/allocated by the caller.
+ O_buffer: [num_pages, bytes_per_page] uint8 — one layer's slab from PackedPagePool
+ (swa page=128 / c4 page=64 / c128 page=2 all share this kernel).
+
+ Per token: 448B fp8(e4m3) in 7x64 ue8m0 groups + 128B bf16 rope in the page data region;
+ 7 exponent bytes (e+127) + 1 zero pad at the page scale tail. Bit-compatible with
+ DeepseekV4MemoryManager._pack_mla_kv + PackedPagePool.write.
+ """
+ seq_len = DestLoc.shape[0]
+ if seq_len == 0:
+ return
+ nope_dim, rope_dim, group_size = 448, 64, 64
+ head_dim = nope_dim + rope_dim
+ scale_bytes = nope_dim // group_size + 1
+
+ KV = KV.reshape(-1, head_dim)
+ assert KV.shape[0] == seq_len, f"Expected KV shape[0]={seq_len}, got {KV.shape[0]}"
+ assert KV.dtype == torch.bfloat16, f"Expected bf16 KV (rope bytes are stored as-is), got {KV.dtype}"
+ bytes_per_page = O_buffer.shape[-1]
+ assert O_buffer.dtype == torch.uint8 and O_buffer.is_contiguous()
+ assert bytes_per_page % 2 == 0
+ assert bytes_per_page >= page_size * (nope_dim + rope_dim * 2 + scale_bytes)
+
+ flat = O_buffer.view(-1)
+ _fwd_kernel_destindex_copy_kv_flashmla_dsv4[(seq_len,)](
+ KV,
+ DestLoc,
+ flat.view(torch.float8_e4m3fn),
+ flat.view(torch.bfloat16),
+ flat,
+ KV.stride(0),
+ KV.stride(1),
+ FP8_MIN=torch.finfo(torch.float8_e4m3fn).min,
+ FP8_MAX=torch.finfo(torch.float8_e4m3fn).max,
+ SCALE_MIN=1e-4,
+ NOPE_DIM=nope_dim,
+ ROPE_DIM=rope_dim,
+ GROUP_SIZE=group_size,
+ NUM_GROUPS=nope_dim // group_size,
+ SCALE_BYTES=scale_bytes,
+ PAGE_SIZE=page_size,
+ BYTES_PER_PAGE=bytes_per_page,
+ num_warps=4,
+ num_stages=1,
+ )
+ return
diff --git a/lightllm/models/deepseek_v4/triton_kernel/gather_c4_indexer_k_dsv4.py b/lightllm/models/deepseek_v4/triton_kernel/gather_c4_indexer_k_dsv4.py
new file mode 100644
index 0000000000..b50eab6726
--- /dev/null
+++ b/lightllm/models/deepseek_v4/triton_kernel/gather_c4_indexer_k_dsv4.py
@@ -0,0 +1,71 @@
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _build_c4_indexer_page_table_kernel(
+ req_idx_ptr, # [batch] int
+ c4_len_ptr, # [batch] int
+ req_to_token_ptr,
+ req_to_token_stride0,
+ full_to_c4_ptr,
+ page_table_ptr, # [batch, page_cap] int32
+ page_cap,
+ hold_req_id,
+ RATIO: tl.constexpr,
+ PAGE_SIZE: tl.constexpr,
+):
+ p = tl.program_id(0)
+ r = tl.program_id(1)
+ req = tl.load(req_idx_ptr + r).to(tl.int64)
+ c4_len = tl.load(c4_len_ptr + r).to(tl.int64)
+ page_start = p * PAGE_SIZE
+ active = (req != hold_req_id) & (page_start < c4_len)
+
+ full_pos0 = page_start * RATIO + (RATIO - 1)
+ full_slot0 = tl.load(
+ req_to_token_ptr + req * req_to_token_stride0 + full_pos0,
+ mask=active,
+ other=0,
+ ).to(tl.int64)
+ c4_slot0 = tl.load(full_to_c4_ptr + full_slot0, mask=active, other=0).to(tl.int64)
+ phys_page = c4_slot0 // PAGE_SIZE
+ tl.store(page_table_ptr + r * page_cap + p, tl.where(active, phys_page, 0).to(tl.int32))
+
+
+@torch.no_grad()
+def build_c4_indexer_page_table(
+ mem_manager,
+ b_req_idx: torch.Tensor,
+ c4_len: torch.Tensor,
+ c4_cap: int,
+ req_to_token_indexs: torch.Tensor,
+ hold_req_id: int,
+):
+ """Build the logical-c4-page -> physical-c4-page table expected by DeepGEMM paged logits.
+
+ Safe only when each logical c4 page maps to a physical page with matching offsets:
+ c4_slot(entry p*64 + o) == page_table[p] * 64 + o
+ which the current token-slot allocator guarantees.
+ """
+ pool = mem_manager.c4_indexer_pool
+ page_size = pool.page_size
+ assert c4_cap % page_size == 0
+ batch = b_req_idx.shape[0]
+ page_cap = c4_cap // page_size
+ page_table = torch.empty((batch, page_cap), dtype=torch.int32, device=b_req_idx.device)
+ _build_c4_indexer_page_table_kernel[(page_cap, batch)](
+ b_req_idx,
+ c4_len,
+ req_to_token_indexs,
+ req_to_token_indexs.stride(0),
+ mem_manager.full_to_c4_indexs,
+ page_table,
+ page_cap,
+ int(hold_req_id),
+ RATIO=4,
+ PAGE_SIZE=page_size,
+ num_warps=1,
+ )
+ return page_table
diff --git a/lightllm/models/deepseek_v4/triton_kernel/norm_rope_cuda.py b/lightllm/models/deepseek_v4/triton_kernel/norm_rope_cuda.py
new file mode 100644
index 0000000000..0ef4434f2e
--- /dev/null
+++ b/lightllm/models/deepseek_v4/triton_kernel/norm_rope_cuda.py
@@ -0,0 +1,98 @@
+import functools
+import hashlib
+import os
+
+import torch
+
+
+@functools.lru_cache(maxsize=1)
+def _load_cuda():
+ from torch.utils.cpp_extension import load
+
+ src = os.path.join(os.path.dirname(__file__), "csrc", "norm_rope.cu")
+ flags = ["-O3"]
+ with open(src, "rb") as source_file:
+ source = source_file.read()
+ capability = torch.cuda.get_device_capability()
+ cache_key = b"\0".join(
+ [
+ source,
+ " ".join(flags).encode(),
+ torch.__version__.encode(),
+ str(torch.version.cuda).encode(),
+ f"sm{capability[0]}{capability[1]}".encode(),
+ os.environ.get("TORCH_CUDA_ARCH_LIST", "").encode(),
+ ]
+ )
+ module_name = f"lightllm_dsv4_norm_rope_v1_{hashlib.sha256(cache_key).hexdigest()[:16]}"
+ return load(
+ name=module_name,
+ sources=[src],
+ extra_cuda_cflags=flags,
+ verbose=False,
+ )
+
+
+def _as_interleaved_freqs(freqs_cis: torch.Tensor) -> torch.Tensor:
+ assert freqs_cis.dtype == torch.complex64
+ freqs_real = torch.view_as_real(freqs_cis).flatten(-2)
+ assert freqs_real.is_contiguous()
+ return freqs_real
+
+
+@torch.no_grad()
+def fused_q_norm_rope(
+ q_input: torch.Tensor,
+ q_output: torch.Tensor,
+ eps: float,
+ freqs_cis: torch.Tensor,
+ positions: torch.Tensor,
+) -> None:
+ _load_cuda().fused_q_norm_rope(q_input, q_output, _as_interleaved_freqs(freqs_cis), positions, float(eps))
+ return
+
+
+@torch.no_grad()
+def fused_k_norm_rope_flashmla(
+ kv: torch.Tensor,
+ kv_weight: torch.Tensor,
+ eps: float,
+ freqs_cis: torch.Tensor,
+ positions: torch.Tensor,
+ out_loc: torch.Tensor,
+ kvcache: torch.Tensor,
+ page_size: int,
+) -> None:
+ _load_cuda().fused_k_norm_rope_flashmla(
+ kv,
+ kv_weight,
+ _as_interleaved_freqs(freqs_cis),
+ positions,
+ out_loc,
+ kvcache,
+ float(eps),
+ int(page_size),
+ )
+ return
+
+
+@torch.no_grad()
+def fused_q_indexer_rope_hadamard_quant(
+ q_input: torch.Tensor,
+ weight: torch.Tensor,
+ weight_scale: float,
+ freqs_cis: torch.Tensor,
+ positions: torch.Tensor,
+):
+ q_fp8 = torch.empty(q_input.shape, dtype=torch.float8_e4m3fn, device=q_input.device)
+ weights_out = torch.empty((*q_input.shape[:-1], 1), dtype=torch.float32, device=q_input.device)
+ _load_cuda().fused_q_indexer_rope_hadamard_quant(
+ q_input,
+ q_fp8,
+ weight,
+ weights_out,
+ float(weight_scale),
+ _as_interleaved_freqs(freqs_cis),
+ positions,
+ )
+ return q_fp8, weights_out
diff --git a/lightllm/models/deepseek_v4/triton_kernel/quant_convert.py b/lightllm/models/deepseek_v4/triton_kernel/quant_convert.py
new file mode 100644
index 0000000000..47d87d4932
--- /dev/null
+++ b/lightllm/models/deepseek_v4/triton_kernel/quant_convert.py
@@ -0,0 +1,16 @@
+import torch
+
+
+def e8m0_to_fp32(scale: torch.Tensor) -> torch.Tensor:
+ """float8_e8m0fnu encodes 2**(byte-127); torch decodes it correctly on .to(float32)."""
+ return scale.to(torch.float32)
+
+
+def dequant_fp8_block_to_bf16(weight_e4m3: torch.Tensor, scale_e8m0: torch.Tensor, block_size: int = 128):
+ """De-quantize an FP8 e4m3 weight [out, in] with block-[bs,bs] ue8m0 scale to bf16."""
+ from lightllm.models.deepseek2.triton_kernel.weight_dequant import weight_dequant
+
+ w = weight_e4m3.cuda().contiguous()
+ s = e8m0_to_fp32(scale_e8m0).cuda().contiguous()
+ # weight_dequant runs with torch default dtype for the output; force bf16 result.
+ return weight_dequant(w, s, block_size)
diff --git a/lightllm/models/deepseek_v4/triton_kernel/topk_transform.py b/lightllm/models/deepseek_v4/triton_kernel/topk_transform.py
new file mode 100644
index 0000000000..76256f5765
--- /dev/null
+++ b/lightllm/models/deepseek_v4/triton_kernel/topk_transform.py
@@ -0,0 +1,61 @@
+import functools
+import hashlib
+import os
+
+import torch
+
+
+@functools.lru_cache(maxsize=1)
+def _load_cuda():
+ from torch.utils.cpp_extension import load
+
+ src = os.path.join(os.path.dirname(__file__), "csrc", "topk_transform.cu")
+ flags = ["-O3"]
+ with open(src, "rb") as source_file:
+ source = source_file.read()
+ capability = torch.cuda.get_device_capability()
+ cache_key = b"\0".join(
+ [
+ source,
+ " ".join(flags).encode(),
+ torch.__version__.encode(),
+ str(torch.version.cuda).encode(),
+ f"sm{capability[0]}{capability[1]}".encode(),
+ os.environ.get("TORCH_CUDA_ARCH_LIST", "").encode(),
+ ]
+ )
+ module_name = f"lightllm_dsv4_topk_v1_{hashlib.sha256(cache_key).hexdigest()[:16]}"
+ return load(
+ name=module_name,
+ sources=[src],
+ extra_cuda_cflags=flags,
+ verbose=False,
+ )
+
+
+@torch.no_grad()
+def topk_transform_512(
+ scores: torch.Tensor,
+ seq_lens: torch.Tensor,
+ page_tables: torch.Tensor,
+ out_page_indices: torch.Tensor,
+ page_size: int,
+ out_raw_indices: torch.Tensor = None,
+) -> None:
+ """Masked top-512 selection over per-token scores + page-translate, for the DeepSeek-V4 c4
+ indexer. Drop-in replacement for the former vendored topk_transform_512 op.
+
+ LightLLM-local CUDA radix port (from SGLang topk_v1.cuh, TVM-FFI stripped and PDL preserved): it
+ early-exits per token at seq_len, so it matches the original perf (unlike torch.topk which
+ scans the full captured c4_cap width). Output is an unordered SET of physical c4 slots (-1 pad).
+
+ Args:
+ scores: [T, c4_cap] fp32 (deep_gemm fp8_paged_mqa_logits output, -inf beyond ctx)
+ seq_lens: [T] int32 (valid causal entries per token)
+ page_tables: [T, npages] int32 (logical->physical c4 page map)
+ out_page_indices: [T, 512] int32 (output physical slots, -1 pad)
+ page_size: c4 pool page size (64)
+ out_raw_indices: optional [T, 512] int32 (raw logical indices, -1 pad)
+ """
+ _load_cuda().topk_transform_512(scores, seq_lens, page_tables, out_page_indices, int(page_size), out_raw_indices)
+ return
diff --git a/lightllm/models/deepseek_v4/workspace.py b/lightllm/models/deepseek_v4/workspace.py
new file mode 100644
index 0000000000..2b260af69a
--- /dev/null
+++ b/lightllm/models/deepseek_v4/workspace.py
@@ -0,0 +1,66 @@
+import torch
+from lightllm.utils.envs_utils import get_env_start_args
+
+
+class DeepseekV4Workspace:
+ def __init__(self, model):
+ self.token_capacity = int(model.batch_max_tokens)
+ self.sliding_window = int(model.config["sliding_window"])
+ self.index_topk = int(model.config["index_topk"])
+ self.c128_cap = self.compress_cap(model.max_seq_length, 128)
+ args = get_env_start_args()
+ overlap = args.enable_decode_microbatch_overlap or args.enable_prefill_microbatch_overlap
+ self.microbatch_count = 1 + int(overlap)
+
+ self.swa_indices = self._alloc(self.sliding_window)
+ self.swa_lengths = torch.empty((self.microbatch_count, self.token_capacity), dtype=torch.int32, device="cuda")
+ self.c4_indices = self._alloc(self.index_topk)
+ self.c4_lengths = torch.empty((self.microbatch_count, self.token_capacity), dtype=torch.int32, device="cuda")
+ self.c128_indices = self._alloc(self.c128_cap)
+ self.c128_lengths = torch.empty((self.microbatch_count, self.token_capacity), dtype=torch.int32, device="cuda")
+ self.flashmla_prefill_q = None
+ self.flashmla_prefill_full_out = None
+
+ def init_flashmla_prefill_q(self, real_q_head_num: int, padded_q_head_num: int, head_dim: int, dtype: torch.dtype):
+ if self.flashmla_prefill_q is None:
+ self.flashmla_prefill_q = torch.empty(
+ (self.token_capacity, padded_q_head_num, head_dim), dtype=dtype, device="cuda"
+ )
+ self.flashmla_prefill_q[:, real_q_head_num:, :].zero_()
+
+ def init_flashmla_prefill_full_out(self, q_head_num: int, head_dim_v: int, dtype: torch.dtype):
+ if self.flashmla_prefill_full_out is None:
+ self.flashmla_prefill_full_out = torch.empty(
+ (self.token_capacity, 1, q_head_num, head_dim_v), dtype=dtype, device="cuda"
+ )
+
+ @staticmethod
+ def compress_cap(max_kv_seq_len: int, ratio: int) -> int:
+ entries = max(1, int(max_kv_seq_len) // ratio)
+ return ((entries + 63) // 64) * 64
+
+ def _alloc(self, width: int) -> torch.Tensor:
+ return torch.empty((self.microbatch_count, self.token_capacity * width), dtype=torch.int32, device="cuda")
+
+ @staticmethod
+ def _view(buffer: torch.Tensor, token_num: int, width: int) -> torch.Tensor:
+ return torch.as_strided(buffer, (token_num, width), (width, 1))
+
+ def swa(self, microbatch_index: int, token_num: int):
+ return (
+ self._view(self.swa_indices[microbatch_index], token_num, self.sliding_window),
+ self.swa_lengths[microbatch_index, :token_num],
+ )
+
+ def c4(self, microbatch_index: int, token_num: int, width: int):
+ assert width <= self.index_topk, f"c4 width {width} exceeds allocated {self.index_topk}"
+ return (
+ self._view(self.c4_indices[microbatch_index], token_num, width),
+ self.c4_lengths[microbatch_index, :token_num],
+ )
+
+ def c128(self, microbatch_index: int, token_num: int, width: int):
+ return (
+ self._view(self.c128_indices[microbatch_index], token_num, width),
+ self.c128_lengths[microbatch_index, :token_num],
+ )
diff --git a/lightllm/models/deepseek_v4_mtp/__init__.py b/lightllm/models/deepseek_v4_mtp/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/lightllm/models/deepseek_v4_mtp/layer_infer/__init__.py b/lightllm/models/deepseek_v4_mtp/layer_infer/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/lightllm/models/deepseek_v4_mtp/layer_infer/pre_layer_infer.py b/lightllm/models/deepseek_v4_mtp/layer_infer/pre_layer_infer.py
new file mode 100644
index 0000000000..66b53ff4aa
--- /dev/null
+++ b/lightllm/models/deepseek_v4_mtp/layer_infer/pre_layer_infer.py
@@ -0,0 +1,54 @@
+import torch
+
+from lightllm.models.deepseek_v4.infer_struct import DeepseekV4InferStateInfo
+from lightllm.models.deepseek_v4_mtp.layer_weights.pre_and_post_layer_weight import (
+ DeepseekV4MTPPreAndPostLayerWeight,
+)
+from lightllm.models.llama.layer_infer.pre_layer_infer import LlamaPreLayerInfer
+
+
+class DeepseekV4MTPPreLayerInfer(LlamaPreLayerInfer):
+ def __init__(self, network_config):
+ super().__init__(network_config)
+ self.eps_ = network_config["rms_norm_eps"]
+ self.hidden_size = network_config["hidden_size"]
+ self.hc_mult = network_config["hc_mult"]
+ return
+
+ def _mtp_forward(
+ self,
+ input_embdings: torch.Tensor,
+ infer_state: DeepseekV4InferStateInfo,
+ layer_weight: DeepseekV4MTPPreAndPostLayerWeight,
+ ):
+ input_embdings = input_embdings.masked_fill(infer_state.position_ids.eq(0).view(-1, 1), 0)
+ target = infer_state.mtp_draft_input_hiddens
+
+ layer_weight.enorm_weight_(input=input_embdings, eps=self.eps_, out=input_embdings)
+ e_proj = layer_weight.e_proj_weight_.mm(input_embdings)
+
+ target = target.view(-1, self.hc_mult, self.hidden_size).contiguous()
+ target = layer_weight.hnorm_weight_(input=target, eps=self.eps_, alloc_func=self.alloc_tensor)
+ h_proj = layer_weight.h_proj_weight_.mm(target.reshape(-1, self.hidden_size))
+ h_proj = h_proj.view(-1, self.hc_mult, self.hidden_size)
+
+ output = h_proj + e_proj.unsqueeze(1)
+ return output.reshape(output.shape[0], self.hc_mult * self.hidden_size)
+
+ def context_forward(
+ self,
+ input_ids,
+ infer_state: DeepseekV4InferStateInfo,
+ layer_weight: DeepseekV4MTPPreAndPostLayerWeight,
+ ):
+ input_embdings = super().context_forward(input_ids, infer_state, layer_weight)
+ return self._mtp_forward(input_embdings, infer_state, layer_weight)
+
+ def token_forward(
+ self,
+ input_ids,
+ infer_state: DeepseekV4InferStateInfo,
+ layer_weight: DeepseekV4MTPPreAndPostLayerWeight,
+ ):
+ input_embdings = super().token_forward(input_ids, infer_state, layer_weight)
+ return self._mtp_forward(input_embdings, infer_state, layer_weight)
diff --git a/lightllm/models/deepseek_v4_mtp/layer_infer/transformer_layer_infer.py b/lightllm/models/deepseek_v4_mtp/layer_infer/transformer_layer_infer.py
new file mode 100644
index 0000000000..651ef9fc24
--- /dev/null
+++ b/lightllm/models/deepseek_v4_mtp/layer_infer/transformer_layer_infer.py
@@ -0,0 +1,8 @@
+from lightllm.models.deepseek_v4.layer_infer.transformer_layer_infer import DeepseekV4TransformerLayerInfer
+
+
+class DeepseekV4MTPTransformerLayerInfer(DeepseekV4TransformerLayerInfer):
+ def __init__(self, layer_num, network_config):
+ super().__init__(layer_num, network_config)
+ self.is_last_layer = True
+ return
diff --git a/lightllm/models/deepseek_v4_mtp/layer_weights/__init__.py b/lightllm/models/deepseek_v4_mtp/layer_weights/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/lightllm/models/deepseek_v4_mtp/layer_weights/pre_and_post_layer_weight.py b/lightllm/models/deepseek_v4_mtp/layer_weights/pre_and_post_layer_weight.py
new file mode 100644
index 0000000000..2b139d016b
--- /dev/null
+++ b/lightllm/models/deepseek_v4_mtp/layer_weights/pre_and_post_layer_weight.py
@@ -0,0 +1,107 @@
+import torch
+
+from lightllm.common.basemodel import PreAndPostLayerWeight
+from lightllm.common.basemodel.layer_weights.meta_weights import (
+ EmbeddingWeight,
+ LMHeadWeight,
+ ParameterWeight,
+ RMSNormWeight,
+ ROWMMWeight,
+)
+from lightllm.common.quantization import Quantcfg
+from lightllm.models.deepseek_v4.triton_kernel.quant_convert import dequant_fp8_block_to_bf16
+
+
+class DeepseekV4MTPPreAndPostLayerWeight(PreAndPostLayerWeight):
+ def __init__(self, data_type, network_config, quant_cfg: Quantcfg):
+ super().__init__(data_type, network_config)
+ self.quant_cfg: Quantcfg = quant_cfg
+
+ hidden = network_config["hidden_size"]
+ vocab = network_config["vocab_size"]
+ hc_mult = network_config["hc_mult"]
+ layer_idx = network_config["n_layer"]
+ prefix = "mtp.0"
+
+ self.wte_weight_ = EmbeddingWeight(
+ dim=hidden,
+ vocab_size=vocab,
+ weight_name=f"{prefix}.emb.tok_emb.weight",
+ data_type=self.data_type_,
+ )
+ self.lm_head_weight_ = LMHeadWeight(
+ dim=hidden,
+ vocab_size=vocab,
+ weight_name=f"{prefix}.head.weight",
+ data_type=self.data_type_,
+ )
+ self.final_norm_weight_ = RMSNormWeight(
+ dim=hidden,
+ weight_name=f"{prefix}.norm.weight",
+ data_type=self.data_type_,
+ )
+
+ self.e_proj_weight_ = ROWMMWeight(
+ in_dim=hidden,
+ out_dims=[hidden],
+ weight_names=f"{prefix}.e_proj.weight",
+ data_type=self.data_type_,
+ quant_method=self.quant_cfg.get_quant_method(layer_idx, "e_proj"),
+ tp_rank=0,
+ tp_world_size=1,
+ )
+ self.h_proj_weight_ = ROWMMWeight(
+ in_dim=hidden,
+ out_dims=[hidden],
+ weight_names=f"{prefix}.h_proj.weight",
+ data_type=self.data_type_,
+ quant_method=self.quant_cfg.get_quant_method(layer_idx, "h_proj"),
+ tp_rank=0,
+ tp_world_size=1,
+ )
+ self.enorm_weight_ = RMSNormWeight(
+ dim=hidden,
+ weight_name=f"{prefix}.enorm.weight",
+ data_type=self.data_type_,
+ )
+ self.hnorm_weight_ = RMSNormWeight(
+ dim=hidden,
+ weight_name=f"{prefix}.hnorm.weight",
+ data_type=self.data_type_,
+ )
+
+ self.hc_head_fn_ = ParameterWeight(
+ weight_name=f"{prefix}.hc_head_fn",
+ data_type=torch.float32,
+ weight_shape=(hc_mult, hc_mult * hidden),
+ )
+ self.hc_head_base_ = ParameterWeight(
+ weight_name=f"{prefix}.hc_head_base",
+ data_type=torch.float32,
+ weight_shape=(hc_mult,),
+ )
+ self.hc_head_scale_ = ParameterWeight(
+ weight_name=f"{prefix}.hc_head_scale",
+ data_type=torch.float32,
+ weight_shape=(1,),
+ )
+ return
+
+ def load_hf_weights(self, weights):
+ self._dequant_in_place(weights)
+ return super().load_hf_weights(weights)
+
+ def _dequant_in_place(self, weights):
+ for attr in (self.e_proj_weight_, self.h_proj_weight_):
+ for weight_name, scale_name in zip(attr.weight_names, attr.weight_scale_names):
+ scale_key = weight_name[: -len(".weight")] + ".scale"
+ if scale_key not in weights:
+ continue
+ if scale_name is None:
+ weights[weight_name] = dequant_fp8_block_to_bf16(weights[weight_name], weights[scale_key]).to(
+ self.data_type_
+ )
+ else:
+ weights[scale_name] = weights[scale_key].to(torch.float32)
+ del weights[scale_key]
+ return
diff --git a/lightllm/models/deepseek_v4_mtp/layer_weights/transformer_layer_weight.py b/lightllm/models/deepseek_v4_mtp/layer_weights/transformer_layer_weight.py
new file mode 100644
index 0000000000..3281bcb98c
--- /dev/null
+++ b/lightllm/models/deepseek_v4_mtp/layer_weights/transformer_layer_weight.py
@@ -0,0 +1,8 @@
+from lightllm.models.deepseek_v4.layer_weights.transformer_layer_weight import DeepseekV4TransformerLayerWeight
+
+
+class DeepseekV4MTPTransformerLayerWeight(DeepseekV4TransformerLayerWeight):
+ def _parse_config(self):
+ super()._parse_config()
+ self.prefix = "mtp.0"
+ return
diff --git a/lightllm/models/deepseek_v4_mtp/model.py b/lightllm/models/deepseek_v4_mtp/model.py
new file mode 100644
index 0000000000..eddcd07d7c
--- /dev/null
+++ b/lightllm/models/deepseek_v4_mtp/model.py
@@ -0,0 +1,142 @@
+import gc
+import os
+from typing import List
+
+import torch
+from safetensors import safe_open
+from tqdm import tqdm
+
+from lightllm.common.basemodel import TpPartBaseModel
+from lightllm.models.deepseek_v4.model import DeepseekV4TpPartModel
+from lightllm.models.deepseek_v4_mtp.layer_infer.pre_layer_infer import DeepseekV4MTPPreLayerInfer
+from lightllm.models.deepseek_v4_mtp.layer_infer.transformer_layer_infer import (
+ DeepseekV4MTPTransformerLayerInfer,
+)
+from lightllm.models.deepseek_v4_mtp.layer_weights.pre_and_post_layer_weight import (
+ DeepseekV4MTPPreAndPostLayerWeight,
+)
+from lightllm.models.deepseek_v4_mtp.layer_weights.transformer_layer_weight import (
+ DeepseekV4MTPTransformerLayerWeight,
+)
+import lightllm.utils.petrel_helper as utils
+from lightllm.utils.log_utils import init_logger
+
+
+logger = init_logger(__name__)
+
+
+class DeepseekV4MTPModel(DeepseekV4TpPartModel):
+ is_mtp_draft_model = True
+
+ pre_and_post_weight_class = DeepseekV4MTPPreAndPostLayerWeight
+ pre_layer_infer_class = DeepseekV4MTPPreLayerInfer
+ transformer_weight_class = DeepseekV4MTPTransformerLayerWeight
+ transformer_layer_infer_class = DeepseekV4MTPTransformerLayerInfer
+
+ def __init__(self, kvargs: dict):
+ self._pre_init(kvargs)
+ super().__init__(kvargs)
+ return
+
+ def _pre_init(self, kvargs: dict):
+ self.main_model: TpPartBaseModel = kvargs.pop("main_model")
+ self.mtp_previous_draft_models: List[TpPartBaseModel] = kvargs.pop("mtp_previous_draft_models")
+ return
+
+ def _init_custom(self):
+ self._freqs_cis_sliding = self.main_model._freqs_cis_sliding
+ self._freqs_cis_compress = self.main_model._freqs_cis_compress
+ self._cos_cached_sliding = self.main_model._cos_cached_sliding
+ self._sin_cached_sliding = self.main_model._sin_cached_sliding
+ self._cos_cached_compress = self.main_model._cos_cached_compress
+ self._sin_cached_compress = self.main_model._sin_cached_compress
+ self.dsv4_workspace = self.main_model.dsv4_workspace
+ for layer in self.layers_infer:
+ layer.freqs_cis = self._freqs_cis_compress if layer.compress_ratio else self._freqs_cis_sliding
+ layer.cos_compress_table = self._cos_cached_compress
+ layer.sin_compress_table = self._sin_cached_compress
+ return
+
+ def _init_req_manager(self):
+ self.req_manager = self.main_model.req_manager
+ return
+
+ def _init_mem_manager(self):
+ self.mem_manager = self.main_model.mem_manager
+ return
+
+ def _init_weights(self, start_layer_index=None):
+ assert start_layer_index is None
+ mtp_layer_index = self.config["n_layer"]
+ self.pre_post_weight = self.pre_and_post_weight_class(
+ self.data_type, network_config=self.config, quant_cfg=self.quant_cfg
+ )
+ self.pre_post_weight.wte_weight_ = self.main_model.pre_post_weight.wte_weight_
+ self.pre_post_weight.lm_head_weight_ = self.main_model.pre_post_weight.lm_head_weight_
+ self.trans_layers_weight = [
+ self.transformer_weight_class(
+ mtp_layer_index,
+ self.data_type,
+ network_config=self.config,
+ quant_cfg=self.quant_cfg,
+ )
+ ]
+ return
+
+ def _init_infer_layer(self, start_layer_index=None):
+ assert start_layer_index is None
+ self.pre_infer = self.pre_layer_infer_class(network_config=self.config)
+ self.post_infer = self.post_layer_infer_class(network_config=self.config)
+ total_pre_layers_num = len(self.main_model.layers_infer)
+ total_pre_layers_num += sum(
+ [len(previous_model.layers_infer) for previous_model in self.mtp_previous_draft_models]
+ )
+ self.layers_infer = [self.transformer_layer_infer_class(total_pre_layers_num, network_config=self.config)]
+ assert self.layers_infer[0].compress_ratio == 0, "DeepSeek-V4 MTP draft layer must be SWA-only"
+ return
+
+ def _init_some_value(self):
+ super()._init_some_value()
+ self.layers_num = 1
+ return
+
+ def _gen_special_model_input(self, token_num: int):
+ return {
+ "mtp_draft_input_hiddens": torch.randn(
+ token_num,
+ self.config["hc_mult"] * self.config["hidden_size"],
+ dtype=self.data_type,
+ device="cuda",
+ )
+ }
+
+ def _load_hf_weights(self):
+ index_file = os.path.join(self.weight_dir_, "model.safetensors.index.json")
+ assert utils.PetrelHelper.exists(index_file), "DeepSeek-V4 MTP requires model.safetensors.index.json."
+ weight_map = utils.PetrelHelper.load_json(index_file)["weight_map"]
+ candidate_files = sorted({file_ for key, file_ in weight_map.items() if key.startswith("mtp.0.")})
+ assert len(candidate_files) > 0, "DeepSeek-V4 MTP weights with prefix mtp.0. were not found."
+
+ loaded_key_count = 0
+ desc = f"pid {os.getpid()} Loading DeepSeek-V4 MTP weights"
+ for file_ in tqdm(candidate_files, total=len(candidate_files), desc=desc):
+ weights = {}
+ with safe_open(os.path.join(self.weight_dir_, file_), "pt", "cpu") as f:
+ for key in f.keys():
+ if key.startswith("mtp.0."):
+ weights[key] = f.get_tensor(key)
+
+ loaded_key_count += len(weights)
+ self.pre_post_weight.load_hf_weights(weights)
+ for layer in self.trans_layers_weight:
+ layer.load_hf_weights(weights)
+ del weights
+ gc.collect()
+
+ self.pre_post_weight.verify_load()
+ [weight.verify_load() for weight in self.trans_layers_weight]
+ logger.info(f"loaded DeepSeek-V4 MTP weights: {loaded_key_count} tensors")
+ return
+
+ def autotune_layers(self):
+ return 1
diff --git a/lightllm/server/api_cli.py b/lightllm/server/api_cli.py
index 1bdf8f3427..70a786edf6 100644
--- a/lightllm/server/api_cli.py
+++ b/lightllm/server/api_cli.py
@@ -161,6 +161,7 @@ def make_argument_parser() -> argparse.ArgumentParser:
"qwen",
"deepseekv31",
"deepseekv32",
+ "deepseekv4",
"glm47",
"kimi_k2",
"qwen3_coder",
@@ -174,6 +175,7 @@ def make_argument_parser() -> argparse.ArgumentParser:
choices=[
"deepseek-r1",
"deepseek-v3",
+ "deepseek-v4",
"glm45",
"gpt-oss",
"kimi",
@@ -620,8 +622,11 @@ def make_argument_parser() -> argparse.ArgumentParser:
type=str,
default=None,
choices=["fp8", "fp4"],
- help="""Expert quantization dtype for EP MoE. Supported values are
- fp8 and fp4. Note that fp4 is only supported on SM100 GPUs.""",
+ help="""Requested dtype for MoE expert weights, fp8 or fp4. Resolves the fused_moe
+ quant method: fp8 -> deepgemm-fp8w8a8-b128; fp4 -> deepgemm-fp4fp8-b32 (online
+ quantization) on SM100 GPUs, or marlin-mxfp4w4a16-b32 (Marlin W4A16, TP only) on other GPUs.
+ Defaults to `expert_dtype` in config.json if present. Per-layer override:
+ --quant_cfg mix_bits with name `fused_moe`.""",
)
parser.add_argument(
"--vit_quant_type",
diff --git a/lightllm/server/api_models.py b/lightllm/server/api_models.py
index 1737d2774d..737a08f628 100644
--- a/lightllm/server/api_models.py
+++ b/lightllm/server/api_models.py
@@ -4,7 +4,8 @@
from pydantic import BaseModel, Field, field_validator, model_validator
from typing import Any, Dict, List, Optional, Union, Literal, ClassVar
-from transformers import GenerationConfig
+
+from lightllm.utils.config_utils import get_generation_config_diff_dict
class ImageURL(BaseModel):
@@ -160,7 +161,7 @@ class CompletionRequest(BaseModel):
def load_generation_cfg(cls, weight_dir: str):
"""Load default values from model generation config."""
try:
- generation_cfg = GenerationConfig.from_pretrained(weight_dir, trust_remote_code=True).to_dict()
+ generation_cfg = get_generation_config_diff_dict(weight_dir)
cls._loaded_defaults = {
"do_sample": generation_cfg.get("do_sample", True),
"presence_penalty": generation_cfg.get("presence_penalty", 0.0),
@@ -221,7 +222,7 @@ class ChatCompletionRequest(BaseModel):
parallel_tool_calls: Optional[bool] = True
# OpenAI parameters for reasoning and others
- reasoning_effort: Optional[Literal["low", "medium", "high"]] = None
+ reasoning_effort: Optional[Literal["none", "low", "medium", "high", "max"]] = None
chat_template_kwargs: Optional[Dict] = None
separate_reasoning: Optional[bool] = True
stream_reasoning: Optional[bool] = False
@@ -242,7 +243,7 @@ class ChatCompletionRequest(BaseModel):
def load_generation_cfg(cls, weight_dir: str):
"""Load default values from model generation config."""
try:
- generation_cfg = GenerationConfig.from_pretrained(weight_dir, trust_remote_code=True).to_dict()
+ generation_cfg = get_generation_config_diff_dict(weight_dir)
cls._loaded_defaults = {
"do_sample": generation_cfg.get("do_sample", True),
"presence_penalty": generation_cfg.get("presence_penalty", 0.0),
diff --git a/lightllm/server/api_openai.py b/lightllm/server/api_openai.py
index 0d934c44c9..1df04af3ab 100644
--- a/lightllm/server/api_openai.py
+++ b/lightllm/server/api_openai.py
@@ -165,8 +165,13 @@ def _is_force_thinking_mode(request: ChatCompletionRequest) -> bool:
return False
if reasoning_parser in ["qwen3-thinking", "gpt-oss", "minimax"]:
return True
- if reasoning_parser in ["deepseek-v3"]:
- return request.chat_template_kwargs is not None and request.chat_template_kwargs.get("thinking") is True
+ if reasoning_parser in ["deepseek-v3", "deepseek-v4"]:
+ chat_template_kwargs = request.chat_template_kwargs or {}
+ if "thinking" in chat_template_kwargs:
+ return chat_template_kwargs["thinking"] is True
+ if request.reasoning_effort is not None:
+ return request.reasoning_effort != "none"
+ return False
if reasoning_parser in ["qwen3", "glm45", "nano_v3", "interns1", "gemma4"]:
# qwen3, glm45, nano_v3, interns1, and gemma4 are reasoning by default;
return not request.chat_template_kwargs or request.chat_template_kwargs.get("enable_thinking", True) is True
diff --git a/lightllm/server/api_start.py b/lightllm/server/api_start.py
index 3cf431d650..40d593d0e8 100644
--- a/lightllm/server/api_start.py
+++ b/lightllm/server/api_start.py
@@ -10,7 +10,11 @@
from .metrics.manager import start_metric_manager
from .embed_cache.manager import start_cache_manager
from lightllm.utils.log_utils import init_logger
-from lightllm.utils.envs_utils import set_env_start_args, set_unique_server_name, get_unique_server_name
+from lightllm.utils.envs_utils import (
+ set_env_start_args,
+ set_unique_server_name,
+ get_unique_server_name,
+)
from lightllm.utils.envs_utils import get_lightllm_gunicorn_keep_alive
from .detokenization.manager import start_detokenization_process
from .router.manager import start_router_process
@@ -333,7 +337,14 @@ def normal_or_p_d_start(args):
from lightllm.utils.config_utils import get_dtype
args.data_type = get_dtype(args.model_dir)
- assert args.data_type in ["fp16", "float16", "bf16", "bfloat16", "fp32", "float32"]
+ assert args.data_type in [
+ "fp16",
+ "float16",
+ "bf16",
+ "bfloat16",
+ "fp32",
+ "float32",
+ ]
already_uesd_ports = [args.port]
if args.nccl_port is not None:
@@ -425,7 +436,6 @@ def normal_or_p_d_start(args):
)
if not args.disable_vision:
-
if not args.visual_use_proxy_mode:
from .visualserver.manager import start_visual_process
@@ -609,7 +619,14 @@ def visual_only_start(args):
from lightllm.utils.config_utils import get_dtype
args.data_type = get_dtype(args.model_dir)
- assert args.data_type in ["fp16", "float16", "bf16", "bfloat16", "fp32", "float32"]
+ assert args.data_type in [
+ "fp16",
+ "float16",
+ "bf16",
+ "bfloat16",
+ "fp32",
+ "float32",
+ ]
logger.info(f"alloced ports: {can_use_ports}")
diff --git a/lightllm/server/build_prompt.py b/lightllm/server/build_prompt.py
index 54d22a0d0d..a28e51dd9e 100644
--- a/lightllm/server/build_prompt.py
+++ b/lightllm/server/build_prompt.py
@@ -53,6 +53,13 @@ def tokenizer_supports_force_thinking() -> bool:
assert tokenizer is not None
+ # Tokenizers that encode prompts in Python (e.g. DeepSeek-V4) have no Jinja
+ # chat_template string to inspect, so advertise thinking support via an
+ # explicit attribute instead.
+ if getattr(tokenizer, "supports_thinking", False):
+ logger.info("tokenizer_supports_force_thinking : True (explicit attribute)")
+ return True
+
try:
ans = "thinking" in tokenizer.chat_template or "enable_thinking" in tokenizer.chat_template
logger.debug(f"chat_template: {tokenizer.chat_template}")
@@ -144,6 +151,8 @@ async def build_prompt(request, tools) -> str:
if request.chat_template_kwargs:
kwargs.update(request.chat_template_kwargs)
+ if request.reasoning_effort is not None and "reasoning_effort" not in kwargs:
+ kwargs["reasoning_effort"] = request.reasoning_effort
# 修复一些parser类型是默认打开thinking，但是 tokenizer有时候不知道打开了thinking。导致
# 构建的reasoning parser 和 tokenizer 的行为不对齐导致的问题。
from .api_openai import _is_force_thinking_mode
@@ -160,5 +169,11 @@ async def build_prompt(request, tools) -> str:
try:
input_str = tokenizer.apply_chat_template(**kwargs, tokenize=False, add_generation_prompt=True, tools=tools)
except Exception as e:
- raise ValueError(f"Failed to build prompt: {e}") from None
+ logger.exception(
+ "Failed to build prompt. request=%s tools=%s template_kwargs=%s",
+ json.dumps(request.model_dump(by_alias=True, exclude_none=True), ensure_ascii=False, default=str),
+ json.dumps(tools, ensure_ascii=False, default=str),
+ json.dumps(kwargs, ensure_ascii=False, default=str),
+ )
+ raise ValueError(f"Failed to build prompt: {e}") from e
return input_str
diff --git a/lightllm/server/core/objs/py_sampling_params.py b/lightllm/server/core/objs/py_sampling_params.py
index cbc63c898d..3e1954502c 100644
--- a/lightllm/server/core/objs/py_sampling_params.py
+++ b/lightllm/server/core/objs/py_sampling_params.py
@@ -4,7 +4,7 @@
"""
import os
from typing import List, Optional, Union, Tuple
-from transformers import GenerationConfig
+from lightllm.utils.config_utils import get_generation_config_diff_dict
from lightllm.server.req_id_generator import MAX_BEST_OF
@@ -110,7 +110,7 @@ def __init__(
@classmethod
def load_generation_cfg(cls, weight_dir):
try:
- generation_cfg = GenerationConfig.from_pretrained(weight_dir, trust_remote_code=True).to_dict()
+ generation_cfg = get_generation_config_diff_dict(weight_dir)
cls._do_sample = generation_cfg.get("do_sample", False)
cls._presence_penalty = generation_cfg.get("presence_penalty", 0.0)
cls._frequency_penalty = generation_cfg.get("frequency_penalty", 0.0)
diff --git a/lightllm/server/core/objs/sampling_params.py b/lightllm/server/core/objs/sampling_params.py
index c39559f5f6..cea422d249 100644
--- a/lightllm/server/core/objs/sampling_params.py
+++ b/lightllm/server/core/objs/sampling_params.py
@@ -1,7 +1,7 @@
import os
import ctypes
from typing import Optional, List, Tuple, Union
-from transformers import GenerationConfig
+from lightllm.utils.config_utils import get_generation_config_diff_dict
from lightllm.server.req_id_generator import MAX_BEST_OF
from .pd_kv_trans_params import PDKVTransParamObj
@@ -395,7 +395,7 @@ def init(self, tokenizer, **kwargs):
@classmethod
def load_generation_cfg(cls, weight_dir):
try:
- generation_cfg = GenerationConfig.from_pretrained(weight_dir, trust_remote_code=True).to_dict()
+ generation_cfg = get_generation_config_diff_dict(weight_dir)
cls._do_sample = generation_cfg.get("do_sample", False)
cls._presence_penalty = generation_cfg.get("presence_penalty", 0.0)
cls._frequency_penalty = generation_cfg.get("frequency_penalty", 0.0)
diff --git a/lightllm/server/core/objs/start_args_type.py b/lightllm/server/core/objs/start_args_type.py
index 40c8028158..8f2df6eba8 100644
--- a/lightllm/server/core/objs/start_args_type.py
+++ b/lightllm/server/core/objs/start_args_type.py
@@ -42,6 +42,7 @@ class StartArgs:
"choices": [
"deepseek-r1",
"deepseek-v3",
+ "deepseek-v4",
"glm45",
"gpt-oss",
"kimi",
diff --git a/lightllm/server/function_call_parser.py b/lightllm/server/function_call_parser.py
index dfcb2f8d9e..a66e383a46 100644
--- a/lightllm/server/function_call_parser.py
+++ b/lightllm/server/function_call_parser.py
@@ -40,6 +40,7 @@
"[TOOL_CALLS]",
"<｜tool▁calls▁begin｜>",
"<｜DSML｜function_calls>",
+ "<｜DSML｜tool_calls>",
]
@@ -1480,11 +1481,14 @@ class DeepSeekV32Detector(BaseFormatDetector):
Reference: https://huggingface.co/deepseek-ai/DeepSeek-V3.2
"""
- def __init__(self):
+ def __init__(self, block_name: str = "function_calls"):
super().__init__()
self.dsml_token = "｜DSML｜"
- self.bot_token = f"<{self.dsml_token}function_calls>"
- self.eot_token = f""
+ # DeepSeek V3.2 wraps tool calls in a `function_calls` block; V4 uses
+ # `tool_calls`. Only the outer block name differs — the invoke/parameter
+ # grammar is identical — so subclasses just override block_name.
+ self.bot_token = f"<{self.dsml_token}{block_name}>"
+ self.eot_token = f""
self.invoke_start_prefix = f"<{self.dsml_token}invoke"
self.invoke_end_token = f""
self.param_end_token = f""
@@ -1568,11 +1572,12 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami
if partial_len:
return StreamingParseResult()
+ normal_text = current_text
self._buffer = ""
for e_token in [self.eot_token, self.invoke_end_token]:
- if e_token in new_text:
- new_text = new_text.replace(e_token, "")
- return StreamingParseResult(normal_text=new_text)
+ if e_token in normal_text:
+ normal_text = normal_text.replace(e_token, "")
+ return StreamingParseResult(normal_text=normal_text)
# Mark that we're inside a function_calls block
if self.has_tool_call(current_text):
@@ -1589,8 +1594,10 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami
try:
# Try to find complete invoke blocks first
- complete_invoke_match = self.invoke_regex.search(current_text)
- if complete_invoke_match:
+ while True:
+ complete_invoke_match = self.invoke_regex.search(current_text)
+ if not complete_invoke_match:
+ break
func_name = complete_invoke_match.group(1)
invoke_body = complete_invoke_match.group(2)
@@ -1654,8 +1661,7 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami
self.current_tool_name_sent = False
self._accumulated_params = []
self.streamed_args_for_tool.append("")
-
- return StreamingParseResult(normal_text="", calls=calls)
+ current_text = self._buffer
# Partial invoke: name is known but parameters are still streaming
partial_match = self.partial_invoke_regex.search(current_text)
@@ -1694,9 +1700,10 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami
if param_matches and len(param_matches) > len(self._accumulated_params):
self._accumulated_params = param_matches
current_args_json = self._dsml_params_to_json(param_matches)
+ open_args_json = current_args_json[:-1] # drop trailing '}'
sent = len(self.streamed_args_for_tool[self.current_tool_id])
- argument_diff = current_args_json[sent:]
+ argument_diff = open_args_json[sent:]
if argument_diff:
calls.append(
@@ -1962,6 +1969,32 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami
self._buffer = current_text[eot_pos + len(self.eot_token) :].lstrip()
+class DeepSeekV4Detector(DeepSeekV32Detector):
+ """
+ Detector for DeepSeek V4 model function call format using DSML.
+
+ Identical grammar to V3.2 (``<｜DSML｜invoke name="...">`` blocks with
+ ``<｜DSML｜parameter name="k" string="true|false">v``
+ tags), except the outer block is named ``tool_calls`` instead of
+ ``function_calls`` — matching the model's own encoding (encoding_dsv4.py:
+ ``tool_calls_block_name = "tool_calls"``) and system prompt.
+
+ Format Structure:
+ ```
+ <｜DSML｜tool_calls>
+ <｜DSML｜invoke name="get_weather">
+ <｜DSML｜parameter name="location" string="true">Hangzhou
+
+
+ ```
+
+ Reference: https://huggingface.co/deepseek-ai/DeepSeek-V4
+ """
+
+ def __init__(self):
+ super().__init__(block_name="tool_calls")
+
+
class FunctionCallParser:
"""
Parser for function/tool calls in model outputs.
@@ -1975,6 +2008,7 @@ class FunctionCallParser:
"deepseekv3": DeepSeekV3Detector,
"deepseekv31": DeepSeekV31Detector,
"deepseekv32": DeepSeekV32Detector,
+ "deepseekv4": DeepSeekV4Detector,
"glm47": Glm47Detector,
"kimi_k2": KimiK2Detector,
"llama3": Llama32Detector,
diff --git a/lightllm/server/reasoning_parser.py b/lightllm/server/reasoning_parser.py
index 8a8d07355b..f351d8a6c8 100644
--- a/lightllm/server/reasoning_parser.py
+++ b/lightllm/server/reasoning_parser.py
@@ -903,6 +903,7 @@ class ReasoningParser:
DetectorMap: Dict[str, Type[BaseReasoningFormatDetector]] = {
"deepseek-r1": DeepSeekR1Detector,
"deepseek-v3": Qwen3Detector,
+ "deepseek-v4": Qwen3Detector,
"glm45": Qwen3Detector,
"gpt-oss": GptOssDetector,
"kimi": KimiDetector,
diff --git a/lightllm/server/router/dynamic_prompt/radix_cache.py b/lightllm/server/router/dynamic_prompt/radix_cache.py
index 21e26c5854..1d58cb8dea 100644
--- a/lightllm/server/router/dynamic_prompt/radix_cache.py
+++ b/lightllm/server/router/dynamic_prompt/radix_cache.py
@@ -2,9 +2,12 @@
import torch
import numpy as np
import collections
-from typing import Tuple, Dict, Set, List, Optional, Union
+from typing import Any, Tuple, Dict, Set, List, Optional, Union
from sortedcontainers import SortedSet
from .shared_arr import SharedArray
+from lightllm.utils.log_utils import init_logger
+
+logger = init_logger(__name__)
class UniqueTimeIdGenerator:
@@ -25,6 +28,7 @@ def __init__(self):
self.parent: TreeNode = None
self.token_id_key: torch.Tensor = None
self.token_mem_index_value: torch.Tensor = None # 用于记录存储的 token_index 为每个元素在 token mem 中的index位置
+ self.token_extra_value: Any = None
self.ref_counter = 0
self.time_id = time_gen.generate_time_id() # 用于标识时间周期
@@ -34,14 +38,17 @@ def __init__(self):
def get_compare_key(self):
return (0 if self.ref_counter == 0 else 1, len(self.children), self.time_id)
- def split_node(self, prefix_len):
+ def split_node(self, prefix_len, child_key_fn=None, extra_value_ops=None):
split_parent_node = TreeNode()
split_parent_node.parent = self.parent
- split_parent_node.parent.children[self.token_id_key[0].item()] = split_parent_node
+ split_parent_node.parent.children[child_key_fn(self.token_id_key)] = split_parent_node
split_parent_node.token_id_key = self.token_id_key[0:prefix_len]
split_parent_node.token_mem_index_value = self.token_mem_index_value[0:prefix_len]
+ if self.token_extra_value is not None and extra_value_ops is not None:
+ split_parent_node.token_extra_value = extra_value_ops.slice(self.token_extra_value, 0, prefix_len)
+ self.token_extra_value = extra_value_ops.slice(self.token_extra_value, prefix_len, len(self.token_id_key))
split_parent_node.children = {}
- split_parent_node.children[self.token_id_key[prefix_len].item()] = self
+ split_parent_node.children[child_key_fn(self.token_id_key[prefix_len:])] = self
split_parent_node.ref_counter = self.ref_counter
new_len = len(split_parent_node.token_mem_index_value)
@@ -56,11 +63,12 @@ def split_node(self, prefix_len):
self.node_prefix_total_len = self.parent.node_prefix_total_len + new_len
return split_parent_node
- def add_and_return_new_child(self, token_id_key, token_mem_index_value):
+ def add_and_return_new_child(self, token_id_key, token_mem_index_value, token_extra_value=None, child_key=None):
child = TreeNode()
child.token_id_key = token_id_key
child.token_mem_index_value = token_mem_index_value
- first_token_key = child.token_id_key[0].item()
+ child.token_extra_value = token_extra_value
+ first_token_key = child.token_id_key[0].item() if child_key is None else child_key
assert first_token_key not in self.children.keys()
self.children[first_token_key] = child
child.parent = self
@@ -71,9 +79,17 @@ def add_and_return_new_child(self, token_id_key, token_mem_index_value):
return child
def remove_child(self, child_node: "TreeNode"):
- del self.children[child_node.token_id_key[0].item()]
- child_node.parent = None
- return
+ child_key = child_node.token_id_key[0].item()
+ if child_key in self.children:
+ del self.children[child_key]
+ child_node.parent = None
+ return
+ for key, value in list(self.children.items()):
+ if value is child_node:
+ del self.children[key]
+ child_node.parent = None
+ return
+ raise KeyError("child node not found")
def update_time(self):
self.time_id = time_gen.generate_time_id()
@@ -103,12 +119,22 @@ class RadixCache:
unique_name 主要用于解决单机，多实列部署时的shm冲突
"""
- def __init__(self, unique_name, total_token_num, rank_in_node, mem_manager=None):
+ def __init__(
+ self,
+ unique_name,
+ total_token_num,
+ rank_in_node,
+ mem_manager=None,
+ page_size: int = 1,
+ extra_value_ops=None,
+ ):
from lightllm.common.kv_cache_mem_manager import MemoryManager
self.mem_manager: MemoryManager = mem_manager
self._key_dtype = torch.int64
self._value_dtype = torch.int64
+ self.page_size = max(1, int(page_size))
+ self.extra_value_ops = extra_value_ops
self.root_node = TreeNode()
self.root_node.token_id_key = torch.zeros((0,), device="cpu", dtype=self._key_dtype)
@@ -124,31 +150,89 @@ def __init__(self, unique_name, total_token_num, rank_in_node, mem_manager=None)
f"{unique_name}_tree_total_tokens_num_{rank_in_node}", (1,), dtype=np.int64
)
self.tree_total_tokens_num.arr[0] = 0
+ self.swa_tree_total_pages_num = 0
+ self.swa_refed_pages_num = 0
+ # 每个 prompt-cache 页折算多少 swa 页(DSV4 为 256/128=2);非 swa 场景为 0,_node_swa_pages_num 退化为常数 0。
+ self._swa_pages_per_prompt_page = self._probe_swa_pages_per_prompt_page()
+
+ def _probe_swa_pages_per_prompt_page(self) -> int:
+ """构造期探测一次 mem_manager 是否带 swa_pool,缓存折算系数,避免热路径反复 getattr。"""
+ if self.mem_manager is None or self.extra_value_ops is None:
+ return 0
+ swa_pool = getattr(self.mem_manager, "swa_pool", None)
+ swa_page_size = getattr(swa_pool, "page_size", None)
+ if swa_page_size is None:
+ return 0
+ return (self.page_size + int(swa_page_size) - 1) // int(swa_page_size)
+
+ def _node_swa_pages_num(self, node: TreeNode) -> int:
+ if self._swa_pages_per_prompt_page == 0 or node.token_extra_value is None:
+ return 0
+ valid = node.token_extra_value.swa_page_valid
+ if valid is None:
+ return 0
+ return int(valid.sum().item()) * self._swa_pages_per_prompt_page
+
+ def _align_len(self, length: int) -> int:
+ if self.page_size <= 1:
+ return int(length)
+ return int(length) // self.page_size * self.page_size
+
+ def align_len(self, length: int) -> int:
+ return self._align_len(length)
+
+ def _child_key(self, key: torch.Tensor):
+ if self.page_size <= 1:
+ return key[0].item()
+ return tuple(key[: self.page_size].tolist())
+
+ def _match_len(self, key: torch.Tensor, node_key: torch.Tensor) -> int:
+ prefix_len = match(key, node_key)
+ return self._align_len(prefix_len)
+
+ def _slice_extra(self, extra_value, start: int, end: int):
+ if extra_value is None:
+ return None
+ assert self.extra_value_ops is not None
+ return self.extra_value_ops.slice(extra_value, start, end)
- def insert(self, key, value=None) -> Tuple[int, Optional[TreeNode]]:
+ def _concat_extra(self, values: list):
+ values = [v for v in values if v is not None]
+ if len(values) == 0:
+ return None
+ assert self.extra_value_ops is not None
+ return self.extra_value_ops.concat(values)
+
+ def insert(self, key, value=None, extra_value=None) -> Tuple[int, Optional[TreeNode]]:
if value is None:
value = key
+ align_len = self._align_len(len(key))
+ key = key[:align_len]
+ value = value[:align_len]
+ if extra_value is not None:
+ extra_value = self._slice_extra(extra_value, 0, align_len)
+
assert len(key) == len(value) # and len(key) >= 1
if len(key) == 0:
return 0, None
- return self._insert_helper(self.root_node, key, value)
+ return self._insert_helper(self.root_node, key, value, extra_value)
- def _insert_helper(self, node: TreeNode, key, value) -> Tuple[int, Optional[TreeNode]]:
+ def _insert_helper(self, node: TreeNode, key, value, extra_value) -> Tuple[int, Optional[TreeNode]]:
handle_stack = collections.deque()
update_list = collections.deque()
- handle_stack.append((node, key, value))
+ handle_stack.append((node, key, value, extra_value))
ans_prefix_len = 0
ans_node = None
while len(handle_stack) != 0:
- node, key, value = handle_stack.popleft()
- ans_tuple = self._insert_helper_no_recursion(node=node, key=key, value=value)
- if len(ans_tuple) == 4:
- (_prefix_len, new_node, new_key, new_value) = ans_tuple
+ node, key, value, extra_value = handle_stack.popleft()
+ ans_tuple = self._insert_helper_no_recursion(node=node, key=key, value=value, extra_value=extra_value)
+ if len(ans_tuple) == 5:
+ (_prefix_len, new_node, new_key, new_value, new_extra_value) = ans_tuple
ans_prefix_len += _prefix_len
- handle_stack.append((new_node, new_key, new_value))
+ handle_stack.append((new_node, new_key, new_value, new_extra_value))
else:
_prefix_len, ans_node = ans_tuple
ans_prefix_len += _prefix_len
@@ -166,15 +250,15 @@ def _insert_helper(self, node: TreeNode, key, value) -> Tuple[int, Optional[Tree
return ans_prefix_len, ans_node
def _insert_helper_no_recursion(
- self, node: TreeNode, key: torch.Tensor, value: torch.Tensor
- ) -> Union[Tuple[int, Optional[TreeNode]], Tuple[int, TreeNode, torch.Tensor, torch.Tensor]]:
+ self, node: TreeNode, key: torch.Tensor, value: torch.Tensor, extra_value=None
+ ) -> Union[Tuple[int, Optional[TreeNode]], Tuple[int, TreeNode, torch.Tensor, torch.Tensor, Any]]:
if node.is_leaf():
self.evict_tree_set.discard(node)
- first_key_id = key[0].item()
+ first_key_id = self._child_key(key)
if first_key_id in node.children.keys():
child: TreeNode = node.children[first_key_id]
- prefix_len = match(key, child.token_id_key)
+ prefix_len = self._match_len(key, child.token_id_key)
if prefix_len == len(key):
if prefix_len == len(child.token_id_key):
if child.is_leaf():
@@ -184,10 +268,14 @@ def _insert_helper_no_recursion(
self.evict_tree_set.add(child)
return prefix_len, child
elif prefix_len < len(child.token_id_key):
+ if prefix_len == 0:
+ return 0, node
if child.is_leaf():
self.evict_tree_set.discard(child)
- split_parent_node = child.split_node(prefix_len)
+ split_parent_node = child.split_node(
+ prefix_len, child_key_fn=self._child_key, extra_value_ops=self.extra_value_ops
+ )
if split_parent_node.is_leaf():
self.evict_tree_set.add(split_parent_node)
@@ -199,15 +287,26 @@ def _insert_helper_no_recursion(
assert False, "can not run to here"
elif prefix_len < len(key) and prefix_len < len(child.token_id_key):
+ if prefix_len == 0:
+ return 0, node
if child.is_leaf():
self.evict_tree_set.discard(child)
+ new_extra_value = self._slice_extra(extra_value, prefix_len, len(key))
key = key[prefix_len:]
value = value[prefix_len:]
- split_parent_node = child.split_node(prefix_len)
- new_node = split_parent_node.add_and_return_new_child(key, value)
+ split_parent_node = child.split_node(
+ prefix_len, child_key_fn=self._child_key, extra_value_ops=self.extra_value_ops
+ )
+ new_node = split_parent_node.add_and_return_new_child(
+ key,
+ value,
+ token_extra_value=new_extra_value,
+ child_key=self._child_key(key),
+ )
# update total token num
self.tree_total_tokens_num.arr[0] += len(new_node.token_mem_index_value)
+ self.swa_tree_total_pages_num += self._node_swa_pages_num(new_node)
if split_parent_node.is_leaf():
self.evict_tree_set.add(split_parent_node)
@@ -218,20 +317,37 @@ def _insert_helper_no_recursion(
self.evict_tree_set.add(child)
return prefix_len, new_node
elif prefix_len < len(key) and prefix_len == len(child.token_id_key):
- return (prefix_len, child, key[prefix_len:], value[prefix_len:])
+ return (
+ prefix_len,
+ child,
+ key[prefix_len:],
+ value[prefix_len:],
+ self._slice_extra(extra_value, prefix_len, len(key)),
+ )
else:
assert False, "can not run to here"
else:
- new_node = node.add_and_return_new_child(key, value)
+ new_node = node.add_and_return_new_child(
+ key,
+ value,
+ token_extra_value=extra_value,
+ child_key=first_key_id,
+ )
# update total token num
self.tree_total_tokens_num.arr[0] += len(new_node.token_mem_index_value)
+ self.swa_tree_total_pages_num += self._node_swa_pages_num(new_node)
if new_node.is_leaf():
self.evict_tree_set.add(new_node)
return 0, new_node
def match_prefix(self, key, update_refs=False):
- assert len(key) != 0
+ key = key[: self._align_len(len(key))]
+ if len(key) == 0:
+ return None, 0, None
+ key = self._trim_key_by_extra_value_validity(key)
+ if len(key) == 0:
+ return None, 0, None
ans_value_list = []
tree_node = self._match_prefix_helper(self.root_node, key, ans_value_list, update_refs=update_refs)
if tree_node != self.root_node:
@@ -245,6 +361,30 @@ def match_prefix(self, key, update_refs=False):
self.dec_node_ref_counter(self.root_node)
return None, 0, None
+ def _trim_key_by_extra_value_validity(self, key: torch.Tensor) -> torch.Tensor:
+ """命中有效性裁剪(extra_value_ops 提供 valid_match_length 时启用,如 DeepSeek-V4 的
+ swa 按页 bitmap): 先做一次只读探测遍历得到自然命中与沿路 extra_value,按其有效边界截短
+ key,随后的正常遍历(加引用/分裂)只走截短后的前缀 —— 引用计数与最终返回值在同一次遍历
+ 内保持一致,不存在事后裁剪导致的漏减/多减。
+
+ 探测遍历可能分裂部分命中的节点(与正常遍历同语义,树不变式不受影响)。裁剪只会缩短命中,
+ 没有任何失败路径。"""
+ if self.extra_value_ops is None:
+ return key
+ valid_match_length = getattr(self.extra_value_ops, "valid_match_length", None)
+ if valid_match_length is None:
+ return key
+ probe_values = []
+ probe_node = self._match_prefix_helper(self.root_node, key, probe_values, update_refs=False)
+ if probe_node == self.root_node or len(probe_values) == 0:
+ return key
+ natural_len = sum(len(v) for v in probe_values)
+ extra_value = self.get_extra_value_by_node(probe_node)
+ valid_len = int(valid_match_length(extra_value, natural_len))
+ if valid_len < natural_len:
+ return key[:valid_len]
+ return key
+
def _match_prefix_helper(
self, node: TreeNode, key: torch.Tensor, ans_value_list: list, update_refs=False
) -> TreeNode:
@@ -286,24 +426,29 @@ def _match_prefix_helper_no_recursion(
# from 0 to 1 need update refs token num
if node.ref_counter == 1:
self.refed_tokens_num.arr[0] += len(node.token_mem_index_value)
+ self.swa_refed_pages_num += self._node_swa_pages_num(node)
if len(key) == 0:
return node
- first_key_id = key[0].item()
+ first_key_id = self._child_key(key)
if first_key_id not in node.children.keys():
return node
else:
child = node.children[first_key_id]
- prefix_len = match(key, child.token_id_key)
+ prefix_len = self._match_len(key, child.token_id_key)
if prefix_len == len(child.token_id_key):
ans_value_list.append(child.token_mem_index_value)
return (child, key[prefix_len:])
elif prefix_len < len(child.token_id_key):
+ if prefix_len == 0:
+ return node
if child.is_leaf():
self.evict_tree_set.discard(child)
- split_parent_node = child.split_node(prefix_len)
+ split_parent_node = child.split_node(
+ prefix_len, child_key_fn=self._child_key, extra_value_ops=self.extra_value_ops
+ )
ans_value_list.append(split_parent_node.token_mem_index_value)
if update_refs:
@@ -311,6 +456,7 @@ def _match_prefix_helper_no_recursion(
# from 0 to 1 need update refs token num
if split_parent_node.ref_counter == 1:
self.refed_tokens_num.arr[0] += len(split_parent_node.token_mem_index_value)
+ self.swa_refed_pages_num += self._node_swa_pages_num(split_parent_node)
if child.is_leaf():
self.evict_tree_set.add(child)
@@ -334,8 +480,11 @@ def evict(self, need_remove_tokens, evict_callback):
), "error evict tree node state"
num_evicted += len(node.token_mem_index_value)
evict_callback(node.token_mem_index_value)
+ if self.extra_value_ops is not None and node.token_extra_value is not None:
+ self.extra_value_ops.free(node.token_extra_value)
# update total token num
self.tree_total_tokens_num.arr[0] -= len(node.token_mem_index_value)
+ self.swa_tree_total_pages_num -= self._node_swa_pages_num(node)
parent_node: TreeNode = node.parent
parent_node.remove_child(node)
if parent_node.is_leaf():
@@ -369,11 +518,12 @@ def _try_merge(self, child_node: TreeNode) -> Optional[TreeNode]:
child_node.token_mem_index_value = torch.cat(
[parent_node.token_mem_index_value, child_node.token_mem_index_value]
)
+ child_node.token_extra_value = self._concat_extra([parent_node.token_extra_value, child_node.token_extra_value])
child_node.node_value_len = len(child_node.token_mem_index_value)
child_node.time_id = max(parent_node.time_id, child_node.time_id)
grandparent_node = parent_node.parent
- key_in_grandparent = parent_node.token_id_key[0].item()
+ key_in_grandparent = self._child_key(parent_node.token_id_key)
grandparent_node.children[key_in_grandparent] = child_node
child_node.parent = grandparent_node
@@ -417,6 +567,8 @@ def clear_tree_nodes(self):
self.tree_total_tokens_num.arr[0] = 0
self.refed_tokens_num.arr[0] = 0
+ self.swa_tree_total_pages_num = 0
+ self.swa_refed_pages_num = 0
return
def dec_node_ref_counter(self, node: TreeNode):
@@ -430,6 +582,7 @@ def dec_node_ref_counter(self, node: TreeNode):
while node is not None:
if node.ref_counter == 1:
self.refed_tokens_num.arr[0] -= len(node.token_mem_index_value)
+ self.swa_refed_pages_num -= self._node_swa_pages_num(node)
node.ref_counter -= 1
node = node.parent
@@ -449,6 +602,7 @@ def add_node_ref_counter(self, node: TreeNode):
while node is not None:
if node.ref_counter == 0:
self.refed_tokens_num.arr[0] += len(node.token_mem_index_value)
+ self.swa_refed_pages_num += self._node_swa_pages_num(node)
node.ref_counter += 1
node = node.parent
@@ -469,12 +623,28 @@ def get_mem_index_value_by_node(self, node: TreeNode) -> Optional[torch.Tensor]:
ans_list.reverse()
return torch.concat(ans_list, dim=0)
+ def get_extra_value_by_node(self, node: TreeNode):
+ if node is None or self.extra_value_ops is None:
+ return None
+
+ ans_list = []
+ while node is not None:
+ if node.token_extra_value is not None:
+ ans_list.append(node.token_extra_value)
+ node = node.parent
+
+ ans_list.reverse()
+ return self._concat_extra(ans_list)
+
def get_refed_tokens_num(self):
return self.refed_tokens_num.arr[0]
def get_tree_total_tokens_num(self):
return self.tree_total_tokens_num.arr[0]
+ def get_unrefed_swa_pages_num(self):
+ return self.swa_tree_total_pages_num - self.swa_refed_pages_num
+
def print_self(self, indent=0):
self._print_helper(self.root_node, indent)
@@ -489,6 +659,63 @@ def _print_helper(self, node: TreeNode, indent):
self._print_helper(child, indent=indent + 2)
return
+ def free_unreferenced_swa_pages(self, need_pages: int) -> None:
+ """DeepSeek-V4 swa free hook: 页 allocator 触底时，回收 ref_count==0 节点的 swa 页。"""
+ if self.mem_manager is None or self.extra_value_ops is None:
+ return
+ allocator = self.mem_manager.swa_page_allocator
+ target = allocator.can_use_mem_size + int(need_pages)
+ evict_slots = []
+ invalidate_payloads = []
+ evict_swa_pages = 0
+ for free_last in (False, True):
+ visited = set()
+ for leaf in self.evict_tree_set:
+ if allocator.can_use_mem_size + evict_swa_pages >= target:
+ break
+ node = leaf
+ while node is not None and node is not self.root_node and node.ref_counter == 0:
+ node_id = id(node)
+ if node_id in visited:
+ node = node.parent
+ continue
+ visited.add(node_id)
+
+ payload = node.token_extra_value
+ if (
+ len(node.token_mem_index_value) > 0
+ and payload is not None
+ and payload.swa_page_valid is not None
+ ):
+ last_page = int(payload.swa_last_valid_page)
+ if last_page >= 0:
+ if free_last:
+ page_slice = slice(last_page, last_page + 1)
+ else:
+ page_slice = slice(0, last_page)
+ valid_pages = int(payload.swa_page_valid[page_slice].sum().item())
+ if valid_pages > 0:
+ start = page_slice.start * self.page_size
+ end = min(page_slice.stop * self.page_size, len(node.token_mem_index_value))
+ if end > start:
+ evict_slots.append(node.token_mem_index_value[start:end])
+ invalidate_payloads.append((payload, page_slice, free_last))
+ evict_swa_pages += valid_pages * self._swa_pages_per_prompt_page
+ if allocator.can_use_mem_size + evict_swa_pages >= target:
+ break
+ node = node.parent
+ if allocator.can_use_mem_size + evict_swa_pages >= target:
+ break
+ if len(evict_slots) == 0:
+ return
+ self.mem_manager.evict_swa(torch.cat(evict_slots))
+ for payload, page_slice, free_last in invalidate_payloads:
+ payload.swa_page_valid[page_slice] = False
+ if free_last:
+ payload.swa_last_valid_page = -1
+ self.swa_tree_total_pages_num -= evict_swa_pages
+ return
+
def free_radix_cache_to_get_enough_token(self, need_token_num):
assert self.mem_manager is not None
if need_token_num > self.mem_manager.allocator.can_use_mem_size:
@@ -504,6 +731,44 @@ def release_mem(mem_index):
self.mem_manager.free(mem_index)
return
+ def _free_radix_full_nodes_until(self, allocator, need: int) -> None:
+ """DeepSeek-V4 压缩池(c4/c128)兑现: 沿 LRU 序逐个驱逐 ref_count==0 的整个 full radix 节点,
+ 经 mem_manager.free() 级联回收其 c4 页 / c128 槽(evict_c4/evict_c128),每驱逐一个就复查
+ *真实* allocator(不靠计数,稳),直到够或已无可驱逐的无引用节点。后者(空闲+可回收仍不足)
+ 由上游 base_backend admission 的 wait_pause 兜底,allocator 的 assert 是最后防线。"""
+ if self.mem_manager is None or allocator is None:
+ return
+ while allocator.can_use_mem_size < need:
+ # 无可驱逐的无引用 token => 停(admission 应已 wait_pause)
+ if self.tree_total_tokens_num.arr[0] <= self.refed_tokens_num.arr[0]:
+ # 兜底没兜住:admission/realize 估算漂移了。打日志便于定位(否则只会撞下游隐晦的
+ # allocator "error alloc state" assert)。
+ logger.warning(
+ f"dsv4 compress-pool realize could not free enough: need={need} "
+ f"free={allocator.can_use_mem_size} tree_total={self.tree_total_tokens_num.arr[0]} "
+ f"refed={self.refed_tokens_num.arr[0]} (admission should have paused this req)"
+ )
+ return
+ release_mems = []
+ # 复用已测的 evict():弹一个 LRU、ref==0 的叶子(>=1 token),其 full 槽经 free 级联回收压缩槽
+ self.evict(1, lambda mem_index: release_mems.append(mem_index))
+ self.mem_manager.free(torch.concat(release_mems))
+ return
+
+ def free_radix_cache_to_get_enough_c4_pages(self, need_pages: int) -> None:
+ allocator = getattr(self.mem_manager, "c4_page_allocator", None) if self.mem_manager is not None else None
+ if allocator is None or need_pages <= 0:
+ return
+ self._free_radix_full_nodes_until(allocator, need_pages)
+ return
+
+ def free_radix_cache_to_get_enough_c128_slots(self, need_slots: int) -> None:
+ allocator = getattr(self.mem_manager, "c128_allocator", None) if self.mem_manager is not None else None
+ if allocator is None or need_slots <= 0:
+ return
+ self._free_radix_full_nodes_until(allocator, need_slots)
+ return
+
class _RadixCacheReadOnlyClient:
"""
diff --git a/lightllm/server/router/model_infer/infer_batch.py b/lightllm/server/router/model_infer/infer_batch.py
index 5c2d0d45fb..6dfc35dbc2 100644
--- a/lightllm/server/router/model_infer/infer_batch.py
+++ b/lightllm/server/router/model_infer/infer_batch.py
@@ -8,7 +8,7 @@
from sortedcontainers import SortedDict
from dataclasses import dataclass, field
from typing import List, Dict, Tuple, Optional, Callable, Any, Union
-from lightllm.common.req_manager import ReqManager, ReqManagerForMamba
+from lightllm.common.req_manager import DeepseekV4ReqManager, ReqManager, ReqManagerForMamba
from lightllm.utils.infer_utils import mark_start, mark_end
from lightllm.server.core.objs import Req, SamplingParams, FinishStatus, ShmReqManager
from lightllm.server.router.dynamic_prompt.radix_cache import RadixCache, TreeNode
@@ -40,6 +40,7 @@ class InferenceContext:
overlap_stream: torch.cuda.Stream = None # 一些情况下推理进程进行异步折叠操作的异步流对象。
cpu_kv_cache_stream: torch.cuda.Stream = None # 用 cpu kv cache 操作的 stream
is_linear_att_mixed_model: bool = False # 标记模型是否是full att 混合 linear att 的混合模型。
+ is_deepseek_v4: bool = False
def register(
self,
@@ -50,7 +51,9 @@ def register(
vocab_size: int,
):
self.args = get_env_start_args()
- from lightllm.server.router.model_infer.mode_backend.base_backend import ModeBackend
+ from lightllm.server.router.model_infer.mode_backend.base_backend import (
+ ModeBackend,
+ )
self.backend: ModeBackend = backend
self.req_manager = req_manager
@@ -64,6 +67,7 @@ def register(
self.vocab_size = vocab_size
self.is_linear_att_mixed_model = isinstance(self.req_manager, ReqManagerForMamba)
+ self.is_deepseek_v4 = isinstance(self.req_manager, DeepseekV4ReqManager)
return
@@ -124,9 +128,17 @@ def add_reqs(self, requests: List[Tuple[int, int, Any, int]], init_prefix_cache:
def free_a_req_mem(self, free_token_index: List, req: "InferReq"):
if self.radix_cache is None:
free_token_index.append(self.req_manager.req_to_token_indexs[req.req_idx][0 : req.cur_kv_len])
+ if self.is_deepseek_v4:
+ # 槽位随 full 槽经 mem_manager.free 级联回收。pause 路径不释放 req_idx,
+ # 必须在此复位出窗水位线 + 清 c128 在途状态(恢复命中走 extend,不会再有
+ # restore/zero 时机;c4 状态随 swa 页生灭,无需处理)。
+ self.req_manager.init_compress_state(req.req_idx)
else:
if not self.is_linear_att_mixed_model:
- self._full_att_free_req(free_token_index=free_token_index, req=req)
+ if self.is_deepseek_v4:
+ self._dsv4_full_att_free_req(free_token_index=free_token_index, req=req)
+ else:
+ self._full_att_free_req(free_token_index=free_token_index, req=req)
else:
self._linear_att_free_req(free_token_index=free_token_index, req=req)
assert len(req.linear_att_len_to_big_page_id) == 0
@@ -149,6 +161,40 @@ def _full_att_free_req(self, free_token_index: List, req: "InferReq"):
req.shared_kv_node = None
return
+ def _dsv4_full_att_free_req(self, free_token_index: List, req: "InferReq"):
+ old_prefix_len = 0 if req.shared_kv_node is None else req.shared_kv_node.node_prefix_total_len
+ inserted_len = old_prefix_len
+ duplicate_prefix_len = old_prefix_len
+
+ cache_len = self.radix_cache.align_len(req.cur_kv_len)
+ self.req_manager: DeepseekV4ReqManager
+ if cache_len > old_prefix_len:
+ payload = self.req_manager.build_prompt_cache_payload(cache_len)
+ value = self.req_manager.req_to_token_indexs[req.req_idx][:cache_len].detach().cpu()
+
+ payload.swa_page_valid = self.req_manager.swa_page_valid_from_watermark(req.req_idx, cache_len)
+ payload.refresh_swa_last_valid_page()
+
+ key = torch.tensor(req.get_input_token_ids()[0:cache_len], dtype=torch.int64, device="cpu")
+ duplicate_prefix_len, _ = self.radix_cache.insert(key, value[:cache_len], extra_value=payload)
+ inserted_len = cache_len
+
+ dense_row = self.req_manager.req_to_token_indexs[req.req_idx]
+ if duplicate_prefix_len > old_prefix_len:
+ free_token_index.append(dense_row[old_prefix_len:duplicate_prefix_len])
+ if req.cur_kv_len > inserted_len:
+ free_token_index.append(dense_row[inserted_len : req.cur_kv_len])
+ if len(free_token_index) == 0:
+ free_token_index.append(dense_row[0:0])
+
+ self.req_manager.init_compress_state(req.req_idx)
+
+ if req.shared_kv_node is not None:
+ assert req.shared_kv_node.node_prefix_total_len <= max(inserted_len, old_prefix_len)
+ self.radix_cache.dec_node_ref_counter(req.shared_kv_node)
+ req.shared_kv_node = None
+ return
+
def _linear_att_free_req(self, free_token_index: List, req: "InferReq"):
assert g_infer_context.is_linear_att_mixed_model is True
args = get_env_start_args()
@@ -325,7 +371,15 @@ def pause_reqs(self, pause_reqs: List["InferReq"], is_master_in_dp: bool):
self.req_manager.free_token(free_token_index)
return self
- def recover_paused_reqs(self, paused_reqs: List["InferReq"], is_master_in_dp: bool, can_alloc_token_num: int):
+ def recover_paused_reqs(
+ self,
+ paused_reqs: List["InferReq"],
+ is_master_in_dp: bool,
+ can_alloc_token_num: int,
+ can_alloc_dsv4_swa_page_num: int = None,
+ can_alloc_dsv4_c4_page_num: int = None,
+ can_alloc_dsv4_c128_slot_num: int = None,
+ ):
if paused_reqs:
for req in paused_reqs:
@@ -333,6 +387,22 @@ def recover_paused_reqs(self, paused_reqs: List["InferReq"], is_master_in_dp: bo
if prefill_need_token_num > can_alloc_token_num:
break
+ swa_page_num = c4_page_num = c128_slot_num = 0
+ if (
+ can_alloc_dsv4_swa_page_num is not None
+ or can_alloc_dsv4_c4_page_num is not None
+ or can_alloc_dsv4_c128_slot_num is not None
+ ):
+ swa_page_num, c4_page_num, c128_slot_num = req.get_dsv4_prefill_need_page_and_slot_num(
+ is_chuncked_prefill=False
+ )
+ if can_alloc_dsv4_swa_page_num is not None and swa_page_num > can_alloc_dsv4_swa_page_num:
+ break
+ if can_alloc_dsv4_c4_page_num is not None and c4_page_num > can_alloc_dsv4_c4_page_num:
+ break
+ if can_alloc_dsv4_c128_slot_num is not None and c128_slot_num > can_alloc_dsv4_c128_slot_num:
+ break
+
if g_infer_context.is_linear_att_mixed_model:
req._linear_match_radix_cache()
else:
@@ -344,6 +414,12 @@ def recover_paused_reqs(self, paused_reqs: List["InferReq"], is_master_in_dp: bo
req.shm_req.is_paused = False
logger.debug(f"infer recover paused req id {req.req_id}")
can_alloc_token_num -= prefill_need_token_num
+ if can_alloc_dsv4_swa_page_num is not None:
+ can_alloc_dsv4_swa_page_num -= swa_page_num
+ if can_alloc_dsv4_c4_page_num is not None:
+ can_alloc_dsv4_c4_page_num -= c4_page_num
+ if can_alloc_dsv4_c128_slot_num is not None:
+ can_alloc_dsv4_c128_slot_num -= c128_slot_num
return
def get_can_alloc_token_num(self):
@@ -354,6 +430,29 @@ def get_can_alloc_token_num(self):
)
return self.req_manager.mem_manager.allocator.can_use_mem_size + radix_cache_unref_token_num
+ def get_can_alloc_dsv4_page_and_slot_num(self):
+ self.req_manager: DeepseekV4ReqManager
+ mem_manager = self.req_manager.mem_manager
+ radix_cache_unref_page_num = 0
+ radix_cache_unref_token_num = 0
+ if self.radix_cache is not None:
+ radix_cache_unref_page_num = self.radix_cache.get_unrefed_swa_pages_num()
+ radix_cache_unref_token_num = (
+ self.radix_cache.get_tree_total_tokens_num() - self.radix_cache.get_refed_tokens_num()
+ )
+ swa_page_num = int(mem_manager.swa_page_allocator.can_use_mem_size) + radix_cache_unref_page_num
+
+ c4_page_num = 0
+ if mem_manager.c4_page_allocator is not None:
+ c4_page_num = int(mem_manager.c4_page_allocator.can_use_mem_size) + int(
+ radix_cache_unref_token_num // self.req_manager.get_prompt_cache_page_size()
+ )
+
+ c128_slot_num = 0
+ if mem_manager.c128_allocator is not None:
+ c128_slot_num = int(mem_manager.c128_allocator.can_use_mem_size) + int(radix_cache_unref_token_num // 128)
+ return swa_page_num, c4_page_num, c128_slot_num
+
def copy_linear_att_state_to_cache_buffer(self, b_req_idx: torch.Tensor, reqs: List["InferReq"]):
"""
该函数用于在线性混合模型prefill后,如果存在大页匹配的情况下，将线性层状态复制到
@@ -382,7 +481,9 @@ def copy_linear_att_state_to_cache_buffer(self, b_req_idx: torch.Tensor, reqs: L
)
big_page_buffer_ids = big_page_buffer_ids.cuda(non_blocking=True)
- from lightllm.common.basemodel.triton_kernel.linear_att_copy import copy_linear_att_state_to_kv_buffer
+ from lightllm.common.basemodel.triton_kernel.linear_att_copy import (
+ copy_linear_att_state_to_kv_buffer,
+ )
copy_linear_att_state_to_kv_buffer(
b_req_idx=b_req_idx,
@@ -412,9 +513,10 @@ def copy_linear_att_state_to_cache_buffer(self, b_req_idx: torch.Tensor, reqs: L
gpu_ssm_state = self.req_manager.req_to_ssm_state.buffer[:, src_buffer_idx, ...]
dst_buffer_idx = req.tail_linear_att_small_page_buffer_id
- dst_conv_state, dst_ssm_state = self.radix_cache.linear_att_small_page_buffers.get_state_cache(
- buffer_idx=dst_buffer_idx
- )
+ (
+ dst_conv_state,
+ dst_ssm_state,
+ ) = self.radix_cache.linear_att_small_page_buffers.get_state_cache(buffer_idx=dst_buffer_idx)
# TODO 对于非连续对象调用 copy_ 效率并不高
dst_conv_state.copy_(gpu_conv_state, non_blocking=True)
dst_ssm_state.copy_(gpu_ssm_state, non_blocking=True)
@@ -563,6 +665,14 @@ def __init__(
self.get_chuncked_input_token_len = self.get_chuncked_input_token_len_for_linear_att
self.get_chuncked_input_token_ids = self.get_chuncked_input_token_ids_for_linear_att
+ if g_infer_context.is_deepseek_v4:
+ mem_manager = g_infer_context.req_manager.mem_manager
+ self.dsv4_swa_page_size: int = mem_manager.swa_pool.page_size
+ self.dsv4_c4_page_size: int = (
+ mem_manager.c4_pool.page_size if mem_manager.c4_page_allocator is not None else 0
+ )
+ self.dsv4_has_c128: bool = mem_manager.c128_allocator is not None
+
self._init_all_state()
self.generator = None
@@ -591,6 +701,8 @@ def _init_all_state(self):
self.cur_output_len = 0
g_infer_context.req_manager.req_sampling_params_manager.init_req_sampling_params(self)
+ if hasattr(g_infer_context.req_manager, "init_compress_state"):
+ g_infer_context.req_manager.init_compress_state(req_idx=self.req_idx)
self.stop_sequences = self.sampling_param.shm_param.stop_sequences.to_list()
# token healing mode 才被使用的管理对象
@@ -626,6 +738,9 @@ def _match_radix_cache(self):
ready_cache_len = share_node.node_prefix_total_len
# 从 cpu 到 gpu 是流内阻塞操作
g_infer_context.req_manager.req_to_token_indexs[self.req_idx, 0:ready_cache_len] = value_tensor
+ # DeepSeek-V4 命中无需任何恢复: 槽位由 full_to_* 映射键控(radix 持有 full 槽即有效,
+ # 命中长度已在 match_prefix 内按 bitmap 裁剪),c4 compressor 状态随 swa 页常驻
+ # (零拷贝续算),c128 状态在 128 对齐边界自然归零(init_compress_state 已清)。
self.cur_kv_len = int(ready_cache_len) # 序列化问题, 该对象可能为numpy.int64，用 int(*)转换
self.shm_req.prompt_cache_len = self.cur_kv_len # 记录 prompt cache 的命中长度
@@ -639,7 +754,10 @@ def _linear_match_radix_cache(self):
enable_prompt_cache = (not self.sampling_param.disable_prompt_cache) and g_infer_context.radix_cache is not None
linear_hash_list = self.shm_req.linear_att_token_hash_list.get_all()
linear_att_hash_page_size = self.args.linear_att_hash_page_size
- match_tokens = min(len(linear_hash_list) * linear_att_hash_page_size, self.get_cur_total_len() - 1)
+ match_tokens = min(
+ len(linear_hash_list) * linear_att_hash_page_size,
+ self.get_cur_total_len() - 1,
+ )
match_tokens = max(0, match_tokens)
match_tokens = (match_tokens // linear_att_hash_page_size) * linear_att_hash_page_size
match_block_num = match_tokens // linear_att_hash_page_size
@@ -705,7 +823,8 @@ def _linear_match_radix_cache(self):
# 将 对应的 value_tensors 中的 kv 数据 拷贝到 tail_mems 中对应的数据去
radix_cache.mem_manager.operator.copy_mem_to_mem(
- value_tensor[cur_big_page_tokens:shared_kv_len], tail_mems
+ value_tensor[cur_big_page_tokens:shared_kv_len],
+ tail_mems,
)
self.shared_kv_node = share_node # 只是为了保证 copy_small_page_buffer_to_linear_att_state 正确调用
@@ -736,7 +855,8 @@ def _linear_match_radix_cache(self):
assert self.tail_linear_att_small_page_buffer_id is None
# 恢复linear att 状态
g_infer_context.req_manager.copy_big_page_buffer_to_linear_att_state(
- big_page_buffer_idx=share_node.big_page_buffer_idx, req=self
+ big_page_buffer_idx=share_node.big_page_buffer_idx,
+ req=self,
)
self.shm_req.shm_cur_kv_len = self.cur_kv_len
@@ -880,6 +1000,53 @@ def _normal_decode_need_token_num(self) -> int:
def _mtp_decode_need_token_num(self) -> int:
return (1 + self.mtp_step) * 2
+ def get_dsv4_prefill_need_page_and_slot_num(self, is_chuncked_prefill: bool) -> Tuple[int, int, int]:
+ start = self.cur_kv_len
+ end = self.get_chuncked_input_token_len() if is_chuncked_prefill else self.get_cur_total_len()
+ if end <= start:
+ return 0, 0, 0
+
+ first_new_page = (start + self.dsv4_swa_page_size - 1) // self.dsv4_swa_page_size
+ last_page = (end - 1) // self.dsv4_swa_page_size
+ swa_page_num = last_page - first_new_page + 1
+
+ c4_page_num = 0
+ first, last = start // 4, end // 4
+ if last > first:
+ # Safe upper bound: touched c4 pages, including a possible already-allocated continuation page.
+ c4_page_num = (last - 1) // self.dsv4_c4_page_size - first // self.dsv4_c4_page_size + 1
+
+ c128_slot_num = max(0, end // 128 - start // 128) if self.dsv4_has_c128 else 0
+ return swa_page_num, c4_page_num, c128_slot_num
+
+ def get_dsv4_decode_need_page_and_slot_num(self) -> Tuple[int, int, int]:
+ seq_len = self.get_cur_total_len()
+ if seq_len <= 0:
+ return 0, 0, 0
+
+ swa_page_num = 0
+ c4_page_num = 0
+ c128_slot_num = 0
+ # Main model prepares current token plus draft-verify rows: SWA + compressed slots.
+ for step in range(self.mtp_step + 1):
+ cur_seq_len = seq_len + step
+ if (cur_seq_len - 1) % self.dsv4_swa_page_size == 0:
+ swa_page_num += 1
+ if cur_seq_len % 4 == 0:
+ entry = cur_seq_len // 4 - 1
+ if entry % self.dsv4_c4_page_size == 0:
+ c4_page_num += 1
+ if self.dsv4_has_c128 and cur_seq_len % 128 == 0:
+ c128_slot_num += 1
+
+ # EAGLE draft forwards after the first one consume newly appended draft-only rows.
+ # The DeepSeek-V4 MTP draft layer is compress_ratio=0, so these rows need only SWA.
+ for step in range(self.mtp_step + 1, self.mtp_step * 2):
+ cur_seq_len = seq_len + step
+ if (cur_seq_len - 1) % self.dsv4_swa_page_size == 0:
+ swa_page_num += 1
+ return swa_page_num, c4_page_num, c128_slot_num
+
class InferReqUpdatePack:
"""
diff --git a/lightllm/server/router/model_infer/mode_backend/base_backend.py b/lightllm/server/router/model_infer/mode_backend/base_backend.py
index a65dfb1bbb..16460bf351 100644
--- a/lightllm/server/router/model_infer/mode_backend/base_backend.py
+++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py
@@ -12,7 +12,7 @@
from lightllm.server.router.model_infer.infer_batch import InferReq, InferReqUpdatePack
from lightllm.server.router.token_load import TokenLoad
from lightllm.common.basemodel.basemodel import TpPartBaseModel
-from lightllm.common.req_manager import ReqManagerForMamba
+from lightllm.common.req_manager import DeepseekV4ReqManager, ReqManagerForMamba
from lightllm.common.linear_att_cache_manager import LinearAttCacheManager
from lightllm.server.router.dynamic_prompt.linear_att_radix_cache import LinearAttPagedRadixCache
from lightllm.server.router.dynamic_prompt.radix_cache import RadixCache
@@ -42,6 +42,7 @@
from lightllm.server.core.objs.shm_objs_io_buffer import ShmObjsIOBuffer
from lightllm.server.router.model_infer.mode_backend.overlap_events import OverlapEventManager, OverlapEventPack
from lightllm.models.deepseek_mtp.model import Deepseek3MTPModel
+from lightllm.models.deepseek_v4_mtp.model import DeepseekV4MTPModel
from lightllm.models.qwen3_moe_mtp.model import Qwen3MOEMTPModel
from lightllm.models.mistral_mtp.model import MistralMTPModel
from lightllm.models.glm4_moe_lite_mtp.model import Glm4MoeLiteMTPModel
@@ -153,6 +154,7 @@ def init_model(self, kvargs):
self.model: TpPartBaseModel = self.model # for easy typing
set_random_seed(2147483647)
self.is_linear_att_mixed_model = isinstance(self.model.req_manager, ReqManagerForMamba)
+ self.is_deepseek_v4 = isinstance(self.model.req_manager, DeepseekV4ReqManager)
if self.is_linear_att_mixed_model:
self.linear_att_cache_manager = LinearAttCacheManager(
@@ -176,12 +178,27 @@ def init_model(self, kvargs):
linear_att_small_page_buffers=self.linear_att_cache_manager,
)
else:
+ radix_page_size = 1
+ radix_extra_value_ops = None
+ if self.is_deepseek_v4:
+ radix_page_size = self.model.req_manager.get_prompt_cache_page_size()
+ radix_extra_value_ops = self.model.req_manager.get_prompt_cache_value_ops()
self.radix_cache = RadixCache(
unique_name=get_unique_server_name(),
total_token_num=self.model.mem_manager.size,
rank_in_node=self.rank_in_node,
mem_manager=self.model.mem_manager,
+ page_size=radix_page_size,
+ extra_value_ops=radix_extra_value_ops,
)
+ if self.is_deepseek_v4:
+ self.model.mem_manager.register_swa_free_hook(self.radix_cache.free_unreferenced_swa_pages)
+
+ if not self.disable_chunked_prefill and radix_page_size > 1:
+ assert self.args.chunked_prefill_size % radix_page_size == 0, (
+ f"chunked_prefill_size={self.args.chunked_prefill_size} must be divisible by "
+ f"prompt-cache page_size={radix_page_size}"
+ )
if "prompt_cache_kv_buffer" in model_cfg:
assert self.use_dynamic_prompt_cache
@@ -333,6 +350,9 @@ def init_mtp_draft_model(self, main_kvargs: dict):
if model_type == "deepseek_v3":
assert self.args.mtp_mode in ["vanilla_with_att", "eagle_with_att"]
self.draft_models.append(Deepseek3MTPModel(mtp_model_kvargs))
+ elif model_type == "deepseek_v4":
+ assert self.args.mtp_mode == "eagle_with_att"
+ self.draft_models.append(DeepseekV4MTPModel(mtp_model_kvargs))
elif model_type == "qwen3_moe":
assert self.args.mtp_mode in ["vanilla_no_att", "eagle_no_att"]
self.draft_models.append(Qwen3MOEMTPModel(mtp_model_kvargs))
@@ -582,6 +602,16 @@ def _get_classed_reqs(
prefill_tokens = 0
can_alloc_token_num = g_infer_context.get_can_alloc_token_num()
+ is_deepseek_v4 = self.is_deepseek_v4
+ can_alloc_dsv4_swa_page_num = None
+ can_alloc_dsv4_c4_page_num = None
+ can_alloc_dsv4_c128_slot_num = None
+ if is_deepseek_v4:
+ (
+ can_alloc_dsv4_swa_page_num,
+ can_alloc_dsv4_c4_page_num,
+ can_alloc_dsv4_c128_slot_num,
+ ) = g_infer_context.get_can_alloc_dsv4_page_and_slot_num()
for req_obj in ready_reqs:
@@ -615,9 +645,21 @@ def _get_classed_reqs(
if is_decode:
token_num = req_obj.decode_need_token_num()
- if token_num <= can_alloc_token_num:
+ can_run = token_num <= can_alloc_token_num
+ if can_run and is_deepseek_v4:
+ swa_page_num, c4_page_num, c128_slot_num = req_obj.get_dsv4_decode_need_page_and_slot_num()
+ can_run = (
+ swa_page_num <= can_alloc_dsv4_swa_page_num
+ and c4_page_num <= can_alloc_dsv4_c4_page_num
+ and c128_slot_num <= can_alloc_dsv4_c128_slot_num
+ )
+ if can_run:
decode_reqs.append(req_obj)
can_alloc_token_num -= token_num
+ if is_deepseek_v4:
+ can_alloc_dsv4_swa_page_num -= swa_page_num
+ can_alloc_dsv4_c4_page_num -= c4_page_num
+ can_alloc_dsv4_c128_slot_num -= c128_slot_num
else:
if wait_pause_count < pause_max_req_num:
req_obj.wait_pause = True
@@ -629,13 +671,28 @@ def _get_classed_reqs(
if req_obj.is_slave_req():
continue
- token_num = req_obj.prefill_need_token_num(is_chuncked_prefill=not self.disable_chunked_prefill)
+ is_chuncked_prefill = not self.disable_chunked_prefill
+ token_num = req_obj.prefill_need_token_num(is_chuncked_prefill=is_chuncked_prefill)
if prefill_tokens + token_num > self.batch_max_tokens:
continue
- if token_num <= can_alloc_token_num:
+ can_run = token_num <= can_alloc_token_num
+ if can_run and is_deepseek_v4:
+ swa_page_num, c4_page_num, c128_slot_num = req_obj.get_dsv4_prefill_need_page_and_slot_num(
+ is_chuncked_prefill=is_chuncked_prefill
+ )
+ can_run = (
+ swa_page_num <= can_alloc_dsv4_swa_page_num
+ and c4_page_num <= can_alloc_dsv4_c4_page_num
+ and c128_slot_num <= can_alloc_dsv4_c128_slot_num
+ )
+ if can_run:
prefill_tokens += token_num
prefill_reqs.append(req_obj)
can_alloc_token_num -= token_num
+ if is_deepseek_v4:
+ can_alloc_dsv4_swa_page_num -= swa_page_num
+ can_alloc_dsv4_c4_page_num -= c4_page_num
+ can_alloc_dsv4_c128_slot_num -= c128_slot_num
else:
if wait_pause_count < pause_max_req_num:
req_obj.wait_pause = True
@@ -655,7 +712,12 @@ def _get_classed_reqs(
if recover_paused:
g_infer_context.recover_paused_reqs(
- paused_reqs=paused_reqs, is_master_in_dp=self.is_master_in_dp, can_alloc_token_num=can_alloc_token_num
+ paused_reqs=paused_reqs,
+ is_master_in_dp=self.is_master_in_dp,
+ can_alloc_token_num=can_alloc_token_num,
+ can_alloc_dsv4_swa_page_num=can_alloc_dsv4_swa_page_num,
+ can_alloc_dsv4_c4_page_num=can_alloc_dsv4_c4_page_num,
+ can_alloc_dsv4_c128_slot_num=can_alloc_dsv4_c128_slot_num,
)
# 在 enable_prefill_decode_mixed 模式下，如果存在 prefill 请求和 decode 请求，
diff --git a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py
index 792a10a788..512bf712fa 100644
--- a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py
+++ b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py
@@ -241,7 +241,7 @@ def decode_mtp(
model_input, run_reqs = prepare_decode_inputs(decode_reqs)
with torch.cuda.stream(g_infer_context.get_overlap_stream()):
- b_mtp_index_cpu = model_input.b_mtp_index
+ b_mtp_index_cpu = model_input.b_mtp_index_cpu
model_output = self.model.forward(model_input)
next_token_ids, next_token_logprobs = sample(model_output.logits, run_reqs, self.eos_id)
# verify the next_token_ids
@@ -345,7 +345,7 @@ def _draft_decode_vanilla(
b_req_mtp_start_loc: torch.Tensor,
):
# share some inference info with the main model
- draft_model_input = main_model_input
+ draft_model_input = main_model_input.make_mtp_draft_input()
draft_model_output = main_model_output
draft_next_token_ids = next_token_ids
all_next_token_ids = []
@@ -387,7 +387,7 @@ def _draft_decode_eagle(
eagle_mem_indexes = eagle_mem_indexes_cpu.cuda(non_blocking=True)
# share some inference info with the main model
- draft_model_input = main_model_input
+ draft_model_input = main_model_input.make_mtp_draft_input()
draft_model_output = main_model_output
draft_next_token_ids = next_token_ids
all_next_token_ids = []
@@ -401,13 +401,13 @@ def _draft_decode_eagle(
draft_model_idx = _step % self.num_mtp_models
draft_model_output: ModelOutput = self.draft_models[draft_model_idx].forward(draft_model_input)
draft_next_token_ids = self._gen_argmax_token_ids(draft_model_output)
- draft_model_input.b_seq_len += 1
- draft_model_input.max_kv_seq_len += 1
+ eagle_mem_indexes_cpu_i = eagle_mem_indexes_cpu[_step * num_reqs : (_step + 1) * num_reqs]
eagle_mem_indexes_i = eagle_mem_indexes[_step * num_reqs : (_step + 1) * num_reqs]
- draft_model_input.mem_indexes = torch.cat(
- [draft_model_input.mem_indexes.view(-1, self.mtp_step + 1)[:, 1:], eagle_mem_indexes_i.view(-1, 1)],
- dim=1,
- ).view(-1)
+ draft_model_input.advance_mtp_decode_step(
+ eagle_mem_indexes_cpu_i,
+ eagle_mem_indexes_i,
+ self.mtp_step,
+ )
all_next_token_ids.append(draft_next_token_ids)
all_next_token_ids = torch.stack(all_next_token_ids, dim=1) # [batch_size, mtp_step + 1]
diff --git a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py b/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py
index e6b9d1c18d..1860fdd5b9 100644
--- a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py
+++ b/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py
@@ -433,7 +433,7 @@ def prefill_mtp(self, event_pack: OverlapEventPack, prefill_reqs: List[InferReq]
def decode_mtp(self, event_pack: OverlapEventPack, decode_reqs: List[InferReq]):
model_input, run_reqs, _ = padded_prepare_decode_inputs(decode_reqs)
- b_mtp_index_cpu = model_input.b_mtp_index
+ b_mtp_index_cpu = model_input.b_mtp_index_cpu
req_num = len(run_reqs)
with torch.cuda.stream(g_infer_context.get_overlap_stream()):
@@ -537,14 +537,13 @@ def _draft_decode_vanilla(
):
all_next_token_ids = []
# share some inference info with the main model
- draft_model_input = model_input
+ draft_model_input = model_input.make_mtp_draft_input()
draft_model_output = model_output
+ all_next_token_ids.append(next_token_ids)
draft_next_token_ids_gpu = torch.zeros((model_input.batch_size), dtype=torch.int64, device="cuda")
if req_num > 0:
draft_next_token_ids_gpu[:req_num].copy_(next_token_ids, non_blocking=True)
- all_next_token_ids.append(draft_next_token_ids_gpu)
-
# process the draft model output
for draft_model_idx in range(self.mtp_step):
@@ -578,7 +577,7 @@ def _draft_decode_eagle(
):
all_next_token_ids = []
# share some inference info with the main model
- draft_model_input = model_input
+ draft_model_input = model_input.make_mtp_draft_input()
draft_model_output = model_output
all_next_token_ids.append(next_token_ids)
draft_next_token_ids_gpu = torch.zeros((model_input.batch_size), dtype=torch.int64, device="cuda")
@@ -587,7 +586,6 @@ def _draft_decode_eagle(
real_req_num = req_num // (self.mtp_step + 1)
padded_req_num = model_input.batch_size // (self.mtp_step + 1) - real_req_num
- eagle_mem_indexes_cpu = None
if g_infer_context.radix_cache is not None:
g_infer_context.radix_cache.free_radix_cache_to_get_enough_token(real_req_num * self.mtp_step)
eagle_mem_indexes_cpu = g_infer_context.req_manager.mem_manager.alloc(real_req_num * self.mtp_step)
@@ -602,19 +600,25 @@ def _draft_decode_eagle(
draft_model_idx = _step % self.num_mtp_models
draft_model_output: ModelOutput = self.draft_models[draft_model_idx].forward(draft_model_input)
# update the meta info of the inference
- draft_model_input.b_seq_len += 1
- draft_model_input.max_kv_seq_len += 1
+ eagle_mem_indexes_cpu_i = eagle_mem_indexes_cpu[_step * real_req_num : (_step + 1) * real_req_num]
eagle_mem_indexes_i = eagle_mem_indexes[_step * real_req_num : (_step + 1) * real_req_num]
+ eagle_mem_indexes_cpu_i = F.pad(
+ input=eagle_mem_indexes_cpu_i,
+ pad=(0, padded_req_num),
+ mode="constant",
+ value=g_infer_context.req_manager.mem_manager.HOLD_TOKEN_MEMINDEX,
+ )
eagle_mem_indexes_i = F.pad(
input=eagle_mem_indexes_i,
pad=(0, padded_req_num),
mode="constant",
value=g_infer_context.req_manager.mem_manager.HOLD_TOKEN_MEMINDEX,
)
- draft_model_input.mem_indexes = torch.cat(
- [draft_model_input.mem_indexes.view(-1, self.mtp_step + 1)[:, 1:], eagle_mem_indexes_i.view(-1, 1)],
- dim=1,
- ).view(-1)
+ draft_model_input.advance_mtp_decode_step(
+ eagle_mem_indexes_cpu_i,
+ eagle_mem_indexes_i,
+ self.mtp_step,
+ )
draft_next_token_ids_gpu = self._gen_argmax_token_ids(draft_model_output)
all_next_token_ids.append(draft_next_token_ids_gpu)
@@ -739,8 +743,8 @@ def decode_overlap_mtp(self, event_pack: OverlapEventPack, decode_reqs: List[Inf
) = padded_overlap_prepare_decode_inputs(decode_reqs)
req_num0, req_num1 = len(run_reqs0), len(run_reqs1)
all_next_token_ids = []
- b_mtp_index_cpu0 = model_input0.b_mtp_index
- b_mtp_index_cpu1 = model_input1.b_mtp_index
+ b_mtp_index_cpu0 = model_input0.b_mtp_index_cpu
+ b_mtp_index_cpu1 = model_input1.b_mtp_index_cpu
with torch.cuda.stream(g_infer_context.get_overlap_stream()):
model_output0, model_output1 = self.model.microbatch_overlap_decode(model_input0, model_input1)
@@ -872,7 +876,8 @@ def _draft_decode_vanilla_overlap(
all_next_token_ids = []
all_next_token_ids.append(next_token_ids)
# share some inference info with the main model
- draft_model_input0, draft_model_input1 = model_input0, model_input1
+ draft_model_input0 = model_input0.make_mtp_draft_input()
+ draft_model_input1 = model_input1.make_mtp_draft_input()
draft_model_output0, draft_model_output1 = model_output0, model_output1
draft_next_token_ids_gpu0 = torch.zeros((model_input0.batch_size), dtype=torch.int64, device="cuda")
@@ -930,7 +935,8 @@ def _draft_decode_eagle_overlap(
all_next_token_ids = []
all_next_token_ids.append(next_token_ids)
# share some inference info with the main model
- draft_model_input0, draft_model_input1 = model_input0, model_input1
+ draft_model_input0 = model_input0.make_mtp_draft_input()
+ draft_model_input1 = model_input1.make_mtp_draft_input()
draft_model_output0, draft_model_output1 = model_output0, model_output1
draft_next_token_ids_gpu0 = torch.zeros((model_input0.batch_size), dtype=torch.int64, device="cuda")
@@ -950,6 +956,8 @@ def _draft_decode_eagle_overlap(
g_infer_context.radix_cache.free_radix_cache_to_get_enough_token(real_req_num * self.mtp_step)
eagle_mem_indexes_cpu = g_infer_context.req_manager.mem_manager.alloc(real_req_num * self.mtp_step)
eagle_mem_indexes = eagle_mem_indexes_cpu.cuda(non_blocking=True)
+ eagle_mem_indexes_cpu0 = eagle_mem_indexes_cpu[0 : real_req_num0 * self.mtp_step]
+ eagle_mem_indexes_cpu1 = eagle_mem_indexes_cpu[real_req_num0 * self.mtp_step : real_req_num * self.mtp_step]
eagle_mem_indexes0 = eagle_mem_indexes[0 : real_req_num0 * self.mtp_step]
eagle_mem_indexes1 = eagle_mem_indexes[real_req_num0 * self.mtp_step : real_req_num * self.mtp_step]
@@ -966,33 +974,45 @@ def _draft_decode_eagle_overlap(
draft_model_input0, draft_model_input1
)
- draft_model_input0.b_seq_len += 1
- draft_model_input0.max_kv_seq_len += 1
+ eagle_mem_indexes_cpu_i = eagle_mem_indexes_cpu0[_step * real_req_num0 : (_step + 1) * real_req_num0]
eagle_mem_indexes_i = eagle_mem_indexes0[_step * real_req_num0 : (_step + 1) * real_req_num0]
+ eagle_mem_indexes_cpu_i = F.pad(
+ input=eagle_mem_indexes_cpu_i,
+ pad=(0, padded_req_num0),
+ mode="constant",
+ value=g_infer_context.req_manager.mem_manager.HOLD_TOKEN_MEMINDEX,
+ )
eagle_mem_indexes_i = F.pad(
input=eagle_mem_indexes_i,
pad=(0, padded_req_num0),
mode="constant",
value=g_infer_context.req_manager.mem_manager.HOLD_TOKEN_MEMINDEX,
)
- draft_model_input0.mem_indexes = torch.cat(
- [draft_model_input0.mem_indexes.view(-1, self.mtp_step + 1)[:, 1:], eagle_mem_indexes_i.view(-1, 1)],
- dim=1,
- ).view(-1)
+ draft_model_input0.advance_mtp_decode_step(
+ eagle_mem_indexes_cpu_i,
+ eagle_mem_indexes_i,
+ self.mtp_step,
+ )
- draft_model_input1.b_seq_len += 1
- draft_model_input1.max_kv_seq_len += 1
+ eagle_mem_indexes_cpu_i = eagle_mem_indexes_cpu1[_step * real_req_num1 : (_step + 1) * real_req_num1]
eagle_mem_indexes_i = eagle_mem_indexes1[_step * real_req_num1 : (_step + 1) * real_req_num1]
+ eagle_mem_indexes_cpu_i = F.pad(
+ input=eagle_mem_indexes_cpu_i,
+ pad=(0, padded_req_num1),
+ mode="constant",
+ value=g_infer_context.req_manager.mem_manager.HOLD_TOKEN_MEMINDEX,
+ )
eagle_mem_indexes_i = F.pad(
input=eagle_mem_indexes_i,
pad=(0, padded_req_num1),
mode="constant",
value=g_infer_context.req_manager.mem_manager.HOLD_TOKEN_MEMINDEX,
)
- draft_model_input1.mem_indexes = torch.cat(
- [draft_model_input1.mem_indexes.view(-1, self.mtp_step + 1)[:, 1:], eagle_mem_indexes_i.view(-1, 1)],
- dim=1,
- ).view(-1)
+ draft_model_input1.advance_mtp_decode_step(
+ eagle_mem_indexes_cpu_i,
+ eagle_mem_indexes_i,
+ self.mtp_step,
+ )
draft_next_token_ids_gpu0 = self._gen_argmax_token_ids(draft_model_output0)
draft_next_token_ids_gpu1 = self._gen_argmax_token_ids(draft_model_output1)
diff --git a/lightllm/server/tokenizer.py b/lightllm/server/tokenizer.py
index e1a4e421d1..6aea7cd672 100644
--- a/lightllm/server/tokenizer.py
+++ b/lightllm/server/tokenizer.py
@@ -90,6 +90,11 @@ def get_tokenizer(
)
logger.info("Using DeepSeek-V3.2 tokenizer mode with Python-based chat template encoding.")
return DeepSeekV32Tokenizer(hf_tokenizer)
+ if model_type == "deepseek_v4":
+ from ..models.deepseek_v4.model import DeepSeekV4Tokenizer
+
+ logger.info("Using DeepSeek-V4 tokenizer mode with Python-based chat template encoding.")
+ return DeepSeekV4Tokenizer(tokenizer, tokenizer_name)
if model_cfg["architectures"][0] == "TarsierForConditionalGeneration":
from ..models.qwen2_vl.vision_process import Qwen2VLImageProcessor
diff --git a/lightllm/utils/config_utils.py b/lightllm/utils/config_utils.py
index c8d7373d54..3c6109829c 100644
--- a/lightllm/utils/config_utils.py
+++ b/lightllm/utils/config_utils.py
@@ -1,6 +1,6 @@
import json
import os
-from typing import Optional, List
+from typing import Any, Dict, Optional, List
from functools import lru_cache
from .envs_utils import get_env_start_args
from lightllm.utils.log_utils import init_logger
@@ -14,6 +14,13 @@ def get_config_json(model_path: str):
return json_obj
+def get_generation_config_diff_dict(model_path: str) -> Dict[str, Any]:
+ from transformers import GenerationConfig
+
+ generation_cfg = GenerationConfig.from_pretrained(model_path, trust_remote_code=True).to_diff_dict()
+ return {key: value for key, value in generation_cfg.items() if value is not None}
+
+
def _derive_max_req_total_len_from_model_config(model_dir: str) -> Optional[int]:
"""
Derive `max_req_total_len` from model config.json.
@@ -464,6 +471,10 @@ def get_tool_call_parser_for_model(model_path: str) -> Optional[str]:
if model_type == "deepseek_v32":
return "deepseekv32"
+ # DeepSeek V4
+ if model_type == "deepseek_v4":
+ return "deepseekv4"
+
return None
@@ -488,8 +499,8 @@ def get_reasoning_parser_for_model(model_path: str) -> Optional[str]:
]:
return "qwen3"
- # DeepSeek V3
- if model_type in ["deepseek_v3", "deepseek_v31", "deepseek_v32"]:
+ # DeepSeek V3 / V4 (share the ... reasoning format, request-gated)
+ if model_type in ["deepseek_v3", "deepseek_v31", "deepseek_v32", "deepseek_v4"]:
return "deepseek-v3"
# DeepSeek R1
diff --git a/test/benchmark/static_inference/static_benchmark.py b/test/benchmark/static_inference/static_benchmark.py
new file mode 100644
index 0000000000..3fd3c62a3e
--- /dev/null
+++ b/test/benchmark/static_inference/static_benchmark.py
@@ -0,0 +1,1542 @@
+"""Static forward benchmark for LightLLM model parts.
+
+The entry uses synthetic token ids and measures forward-only TPS for prefill,
+chunked prefill, decode, and MTP decode cases.
+"""
+
+import argparse
+import json
+import math
+import os
+import queue
+import sys
+import time
+import traceback
+from dataclasses import asdict, dataclass, replace
+from pathlib import Path
+from types import SimpleNamespace
+from typing import Dict, List, Optional, Sequence
+
+import numpy as np
+import torch
+import torch.multiprocessing as mp
+from transformers import PretrainedConfig
+
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+if str(REPO_ROOT) not in sys.path:
+ sys.path.append(str(REPO_ROOT))
+
+from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput
+from lightllm.models import get_model
+from lightllm.models.deepseek_mtp.model import Deepseek3MTPModel
+from lightllm.models.deepseek_v4_mtp.model import DeepseekV4MTPModel
+from lightllm.models.glm4_moe_lite_mtp.model import Glm4MoeLiteMTPModel
+from lightllm.models.mistral_mtp.model import MistralMTPModel
+from lightllm.models.qwen3_moe_mtp.model import Qwen3MOEMTPModel
+from lightllm.server.api_cli import make_argument_parser
+from lightllm.utils.config_utils import get_dtype, get_vocab_size
+from lightllm.utils.dist_utils import init_distributed_env
+from lightllm.utils.envs_utils import set_env_start_args
+
+
+DEFAULT_BATCH_SIZES = [2, 8, 16, 32, 64, 128]
+MTP_MODES = {"vanilla_with_att", "eagle_with_att", "vanilla_no_att", "eagle_no_att"}
+PREFILL_TABLE_HEADERS = [
+ "ctx",
+ "hit",
+ "bs",
+ "max_total_token_num",
+ "uncached",
+ "cached",
+ "tokens",
+ "ms",
+ "qps",
+ "tok/s",
+ "logical_tok/s",
+]
+DECODE_TABLE_HEADERS = [
+ "ctx",
+ "bs",
+ "accept",
+ "max_total_token_num",
+ "ms",
+ "qps",
+ "tok/s",
+ "itl_ms",
+]
+
+
+@dataclass(frozen=True)
+class BenchmarkCase:
+ name: str
+ stage: str
+ batch_size: int
+ context_len: int
+ output_len: int
+ chunked_prefill_size: Optional[int] = None
+ profiled_max_total_token_num: Optional[int] = None
+ profiled_batch_divisor: Optional[int] = None
+ cache_hit_rate: float = 0.0
+ prefill_uncached_len: Optional[int] = None
+ prefill_step_tokens_per_req: Optional[int] = None
+ prefill_batch_size_by_batch_max_tokens: Optional[int] = None
+
+
+@dataclass
+class BenchmarkResult:
+ case: str
+ stage: str
+ batch_size: int
+ context_len: int
+ output_len: int
+ chunked_prefill_size: Optional[int]
+ elapsed_ms: float
+ measured_tokens: int
+ qps: float
+ tps: float
+ profiled_max_total_token_num: Optional[int] = None
+ profiled_batch_divisor: Optional[int] = None
+ ttft_ms: Optional[float] = None
+ inter_token_latency_ms: Optional[float] = None
+ cache_hit_rate: float = 0.0
+ prefill_uncached_len: Optional[int] = None
+ prefill_cached_len: Optional[int] = None
+ prefill_step_tokens_per_req: Optional[int] = None
+ mtp_accept_rate: Optional[float] = None
+ logical_tps: Optional[float] = None
+
+
+class TokenSource:
+ def __init__(self, args: SimpleNamespace):
+ self.vocab_size = max(2, int(get_vocab_size(args.model_dir) or 0))
+ self.rng = np.random.default_rng(args.seed)
+
+ def batch(self, batch_size: int, need_len: int) -> np.ndarray:
+ return self.rng.integers(low=0, high=self.vocab_size, size=(batch_size, need_len), dtype=np.int64)
+
+
+def cpu_i32_full(shape, value) -> torch.Tensor:
+ return torch.full(shape, value, dtype=torch.int32, device="cpu")
+
+
+def cpu_i32_zeros(size: int) -> torch.Tensor:
+ return torch.zeros(size, dtype=torch.int32, device="cpu")
+
+
+def empty_multimodal_params(batch_size: int) -> List[Dict]:
+ return [{"images": [], "audios": []} for _ in range(batch_size)]
+
+
+class StaticBenchmarkExecutor:
+ def __init__(
+ self,
+ args: SimpleNamespace,
+ model,
+ draft_models: List,
+ token_source: TokenSource,
+ ):
+ self.args = args
+ self.model = model
+ self.draft_models = draft_models
+ self.token_source = token_source
+
+ def _case_iters(self, warmup: bool) -> int:
+ return self.args.warmup_iters if warmup else self.args.bench_iters
+
+ def run_case(self, case: BenchmarkCase, warmup: bool) -> BenchmarkResult:
+ if case.stage == "prefill":
+ return self._run_prefill_case(case, warmup)
+ if case.stage == "decode":
+ return self._run_decode_case(case, warmup)
+ raise ValueError(f"unknown benchmark stage: {case.stage}")
+
+ def _run_prefill_case(self, case: BenchmarkCase, warmup: bool) -> BenchmarkResult:
+ """Measure full uncached prefill, chunked by production admission size."""
+ uncached_len = int(case.prefill_uncached_len or case.context_len)
+ cached_len = max(0, case.context_len - uncached_len)
+ tokens = self.token_source.batch(case.batch_size, uncached_len)
+ elapsed = 0.0
+ measured_tokens = case.batch_size * uncached_len
+
+ for _ in range(self._case_iters(warmup)):
+ self._reset_model_cache()
+ req_idx = self._alloc_req_indexes(case.batch_size)
+ if cached_len > 0:
+ self._materialize_cached_prefix(req_idx, cached_len)
+ inputs = self._build_prefill_inputs(
+ token_rows=tokens,
+ req_idx=req_idx,
+ prompt_len=uncached_len,
+ chunk_size=case.chunked_prefill_size,
+ initial_ready_cache_len=cached_len,
+ )
+ torch.cuda.synchronize()
+ start = time.perf_counter()
+ output = None
+ for model_input in inputs:
+ output = self._forward_prefill_input(model_input, allow_overlap=True)
+ torch.cuda.synchronize()
+ elapsed += time.perf_counter() - start
+ self._touch_output(output)
+
+ self._reset_model_cache()
+ return self._make_result(case, elapsed, measured_tokens, warmup)
+
+ def _run_decode_case(self, case: BenchmarkCase, warmup: bool) -> BenchmarkResult:
+ mtp_enabled = self._mtp_enabled()
+ measured_tokens = case.batch_size * case.output_len
+ elapsed = 0.0
+ decode_step_count = 0
+ iters = self._case_iters(warmup)
+
+ for _ in range(iters):
+ self._reset_model_cache()
+ req_idx, seq_len, next_ids = self._materialize_context_for_decode(case)
+ if mtp_enabled:
+ step_elapsed, step_count = self._run_mtp_decode_steps(
+ case=case,
+ req_idx=req_idx,
+ seq_len=seq_len,
+ next_ids=next_ids,
+ )
+ elapsed += step_elapsed
+ decode_step_count += step_count
+ else:
+ elapsed += self._run_plain_decode_steps(
+ case=case,
+ req_idx=req_idx,
+ seq_len=seq_len,
+ next_ids=next_ids,
+ )
+ decode_step_count += case.output_len
+
+ self._reset_model_cache()
+ inter_token_latency_ms = elapsed * 1000.0 / max(1, decode_step_count) if iters > 0 else None
+ return self._make_result(
+ case,
+ elapsed,
+ measured_tokens,
+ warmup,
+ inter_token_latency_ms=inter_token_latency_ms,
+ )
+
+ def _materialize_context_for_decode(self, case: BenchmarkCase):
+ """Allocate historical KV slots so decode can be measured without prefill."""
+ req_idx = self._alloc_req_indexes(case.batch_size)
+ self._materialize_cached_prefix(req_idx, case.context_len)
+ seq_len = cpu_i32_full((case.batch_size,), case.context_len)
+ next_ids = torch.from_numpy(np.ascontiguousarray(self.token_source.batch(case.batch_size, 1).reshape(-1))).to(
+ torch.int64
+ )
+ return req_idx, seq_len, next_ids
+
+ def _run_plain_decode_steps(
+ self,
+ case: BenchmarkCase,
+ req_idx: torch.Tensor,
+ seq_len: torch.Tensor,
+ next_ids: torch.Tensor,
+ ) -> float:
+ elapsed = 0.0
+ for step in range(case.output_len):
+ seq_len += 1
+ model_input = self._make_decode_input(
+ batch_size=case.batch_size,
+ req_idx=req_idx,
+ mtp_index=cpu_i32_zeros(case.batch_size),
+ seq_len=seq_len,
+ input_ids=next_ids.reshape(-1),
+ max_kv_seq_len=int(seq_len.max().item()),
+ mem_token_num=case.batch_size,
+ )
+ torch.cuda.synchronize()
+ start = time.perf_counter()
+ output = self._forward_decode_input(model_input, allow_overlap=True)
+ torch.cuda.synchronize()
+ elapsed += time.perf_counter() - start
+ self._touch_output(output)
+
+ next_ids = self._argmax_ids(output.logits)
+ return elapsed
+
+ def _run_mtp_decode_steps(
+ self,
+ case: BenchmarkCase,
+ req_idx: torch.Tensor,
+ seq_len: torch.Tensor,
+ next_ids: torch.Tensor,
+ ) -> tuple:
+ elapsed = 0.0
+ step_count = 0
+ generated_len = 0
+ step_width = self._mtp_step_width()
+ base_req_idx, b_mtp_index = self._build_mtp_decode_index_tensors(req_idx, step_width)
+ current_candidates = next_ids
+
+ while generated_len < case.output_len:
+ accepted_width = self._sample_mtp_accept_width(step_width, case.output_len - generated_len)
+ if current_candidates.ndim == 1:
+ current_candidates = current_candidates[:, None].repeat(1, step_width)
+
+ b_seq_len = self._build_mtp_seq_len(seq_len, step_width)
+ model_input = self._make_decode_input(
+ batch_size=case.batch_size * step_width,
+ req_idx=base_req_idx,
+ mtp_index=b_mtp_index,
+ seq_len=b_seq_len,
+ input_ids=current_candidates.reshape(-1),
+ max_kv_seq_len=int(b_seq_len.max().item()),
+ mem_token_num=case.batch_size * step_width,
+ )
+
+ torch.cuda.synchronize()
+ start = time.perf_counter()
+ output = self.model.forward(model_input)
+ candidate_rows, temporary_mem = self._run_mtp_draft_decode(
+ model_input=model_input,
+ model_output=output,
+ real_batch_size=case.batch_size,
+ step_width=step_width,
+ )
+ torch.cuda.synchronize()
+ elapsed += time.perf_counter() - start
+ self._touch_output(output)
+ if temporary_mem is not None:
+ self.model.req_manager.mem_manager.free(temporary_mem)
+
+ self._free_rejected_mtp_mem(
+ model_input=model_input,
+ real_batch_size=case.batch_size,
+ step_width=step_width,
+ accepted_width=accepted_width,
+ )
+ current_candidates = (
+ self._select_mtp_candidates(
+ candidate_rows=candidate_rows,
+ real_batch_size=case.batch_size,
+ step_width=step_width,
+ accepted_width=accepted_width,
+ )
+ .detach()
+ .cpu()
+ )
+ seq_len += accepted_width
+ generated_len += accepted_width
+ step_count += 1
+
+ return elapsed, step_count
+
+ def _run_mtp_draft_decode(
+ self,
+ model_input: ModelInput,
+ model_output: ModelOutput,
+ real_batch_size: int,
+ step_width: int,
+ ):
+ draft_input = model_input.make_mtp_draft_input()
+ draft_output = model_output
+ draft_next_ids = self._argmax_ids(model_output.logits).cuda(non_blocking=True)
+ generated = [draft_next_ids.detach()]
+
+ temporary_mem = None
+ if self.args.mtp_mode.startswith("eagle"):
+ temporary_mem = self.model.req_manager.mem_manager.alloc(real_batch_size * self.args.mtp_step)
+ temporary_mem_gpu = temporary_mem.cuda(non_blocking=True)
+ else:
+ temporary_mem_gpu = None
+
+ for step in range(self.args.mtp_step):
+ draft_input.input_ids = draft_next_ids
+ draft_input.mtp_draft_input_hiddens = draft_output.mtp_main_output_hiddens
+ draft_model = self.draft_models[step % self._num_mtp_modules()]
+ draft_output = draft_model.forward(draft_input)
+ draft_next_ids = self._argmax_ids(draft_output.logits).cuda(non_blocking=True)
+ generated.append(draft_next_ids.detach())
+
+ if self.args.mtp_mode.startswith("eagle"):
+ mem_i_cpu = temporary_mem[step * real_batch_size : (step + 1) * real_batch_size]
+ mem_i = temporary_mem_gpu[step * real_batch_size : (step + 1) * real_batch_size]
+ draft_input.advance_mtp_decode_step(mem_i_cpu, mem_i, self.args.mtp_step)
+
+ return torch.stack(generated[:step_width], dim=1), temporary_mem
+
+ def _sample_mtp_accept_width(self, step_width: int, remaining_tokens: int) -> int:
+ """Sample accepted MTP width outside the timed decode section."""
+ accept_rate = float(self.args.mtp_accept_rate)
+ accepted_width = 1
+ for _ in range(step_width - 1):
+ if self.token_source.rng.random() >= accept_rate:
+ break
+ accepted_width += 1
+ return max(1, min(accepted_width, remaining_tokens))
+
+ def _select_mtp_candidates(
+ self,
+ candidate_rows: torch.Tensor,
+ real_batch_size: int,
+ step_width: int,
+ accepted_width: int,
+ ) -> torch.Tensor:
+ row_ids = torch.arange(real_batch_size, device=candidate_rows.device) * step_width + accepted_width - 1
+ return candidate_rows.index_select(0, row_ids)
+
+ def _free_rejected_mtp_mem(
+ self,
+ model_input: ModelInput,
+ real_batch_size: int,
+ step_width: int,
+ accepted_width: int,
+ ):
+ if accepted_width >= step_width:
+ return
+ rejected_mem = (
+ model_input.mem_indexes_cpu.view(real_batch_size, step_width)[:, accepted_width:].contiguous().reshape(-1)
+ )
+ if rejected_mem.numel() > 0:
+ self.model.req_manager.mem_manager.free(rejected_mem)
+
+ def _build_prefill_inputs(
+ self,
+ token_rows: np.ndarray,
+ req_idx: torch.Tensor,
+ prompt_len: int,
+ chunk_size: Optional[int],
+ initial_ready_cache_len: int = 0,
+ ) -> List[ModelInput]:
+ if not chunk_size or chunk_size <= 0 or chunk_size >= prompt_len:
+ return [
+ self._make_prefill_input(
+ token_rows[:, :prompt_len],
+ req_idx,
+ ready_cache_len=initial_ready_cache_len,
+ )
+ ]
+
+ inputs = []
+ for start in range(0, prompt_len, chunk_size):
+ end = min(prompt_len, start + chunk_size)
+ inputs.append(
+ self._make_prefill_input(
+ token_rows[:, start:end],
+ req_idx,
+ ready_cache_len=initial_ready_cache_len + start,
+ )
+ )
+ return inputs
+
+ def _materialize_cached_prefix(self, req_idx: torch.Tensor, cached_len: int):
+ """Allocate dummy prefix KV so cache-hit cases consume real capacity."""
+ if cached_len <= 0:
+ return
+ batch_size = int(req_idx.shape[0])
+ need_tokens = batch_size * cached_len
+ mem_indexes = self.model.req_manager.mem_manager.alloc(need_tokens)
+ if mem_indexes is None:
+ raise RuntimeError(f"failed to allocate cached prefix: bs={batch_size} cached_len={cached_len}")
+ req_idx_gpu = req_idx.cuda(non_blocking=True)
+ mem_indexes_gpu = mem_indexes.reshape(batch_size, cached_len).cuda(non_blocking=True)
+ self.model.req_manager.req_to_token_indexs[req_idx_gpu, :cached_len] = mem_indexes_gpu
+ self._materialize_cached_prefix_extra_slots(req_idx, cached_len)
+
+ def _materialize_cached_prefix_extra_slots(self, req_idx: torch.Tensor, cached_len: int):
+ req_manager = self.model.req_manager
+ batch_size = int(req_idx.shape[0])
+ b_req_idx = req_idx.cuda(non_blocking=True)
+ b_seq_len_cpu = cpu_i32_full((batch_size,), cached_len)
+ b_seq_len = b_seq_len_cpu.cuda(non_blocking=True)
+
+ if hasattr(req_manager, "prepare_prefill_compress_slots"):
+ b_ready_cache_len_cpu = cpu_i32_zeros(batch_size)
+ req_manager.prepare_prefill_compress_slots(
+ b_req_idx=b_req_idx,
+ b_ready_cache_len=b_ready_cache_len_cpu.cuda(non_blocking=True),
+ b_seq_len=b_seq_len,
+ b_req_idx_cpu=req_idx,
+ b_ready_cache_len_cpu=b_ready_cache_len_cpu,
+ b_seq_len_cpu=b_seq_len_cpu,
+ )
+
+ if hasattr(req_manager, "prepare_prefill_swa"):
+ swa_ready_len = self._cached_prefix_swa_ready_len(cached_len)
+ b_ready_cache_len_cpu = cpu_i32_full((batch_size,), swa_ready_len)
+ req_manager.prepare_prefill_swa(
+ b_req_idx=b_req_idx,
+ b_ready_cache_len=b_ready_cache_len_cpu.cuda(non_blocking=True),
+ b_seq_len=b_seq_len,
+ b_req_idx_cpu=req_idx,
+ b_ready_cache_len_cpu=b_ready_cache_len_cpu,
+ b_seq_len_cpu=b_seq_len_cpu,
+ )
+
+ def _cached_prefix_swa_ready_len(self, cached_len: int) -> int:
+ req_manager = self.model.req_manager
+ retain_fn = getattr(req_manager, "_swa_retain_len", None)
+ if retain_fn is not None:
+ retain_len = int(retain_fn())
+ else:
+ retain_len = int(getattr(req_manager, "sliding_window", cached_len) or cached_len)
+ ready_len = max(0, int(cached_len) - max(1, retain_len))
+ page_fn = getattr(req_manager, "get_prompt_cache_page_size", None)
+ page_size = int(page_fn()) if page_fn is not None else 1
+ return ready_len // max(1, page_size) * max(1, page_size)
+
+ def _make_prefill_input(self, token_chunk: np.ndarray, req_idx: torch.Tensor, ready_cache_len: int) -> ModelInput:
+ batch_size, q_len = token_chunk.shape
+ seq_len_value = ready_cache_len + q_len
+ b_seq_len = cpu_i32_full((batch_size,), seq_len_value)
+ b_ready_cache_len = cpu_i32_full((batch_size,), ready_cache_len)
+ b_q_seq_len = b_seq_len - b_ready_cache_len
+ b_prefill_start_loc = b_q_seq_len.cumsum(dim=0, dtype=torch.int32) - b_q_seq_len
+ input_ids = torch.from_numpy(np.ascontiguousarray(token_chunk.reshape(-1))).to(torch.int64)
+ mem_indexes = self.model.req_manager.mem_manager.alloc(input_ids.shape[0])
+ return ModelInput(
+ batch_size=batch_size,
+ total_token_num=int(b_seq_len.sum().item()),
+ max_q_seq_len=q_len,
+ max_kv_seq_len=seq_len_value,
+ max_cache_len=ready_cache_len,
+ prefix_total_token_num=ready_cache_len * batch_size,
+ input_ids=input_ids,
+ b_req_idx=req_idx,
+ b_mtp_index=cpu_i32_zeros(batch_size),
+ b_seq_len=b_seq_len,
+ mem_indexes_cpu=mem_indexes,
+ is_prefill=True,
+ b_ready_cache_len=b_ready_cache_len,
+ b_prefill_start_loc=b_prefill_start_loc,
+ b_prefill_has_output_cpu=[False] * batch_size,
+ multimodal_params=empty_multimodal_params(batch_size),
+ )
+
+ def _make_decode_input(
+ self,
+ batch_size: int,
+ req_idx: torch.Tensor,
+ mtp_index: torch.Tensor,
+ seq_len: torch.Tensor,
+ input_ids: torch.Tensor,
+ max_kv_seq_len: int,
+ mem_token_num: int,
+ ) -> ModelInput:
+ mem_indexes = self.model.req_manager.mem_manager.alloc(mem_token_num)
+ return ModelInput(
+ batch_size=batch_size,
+ total_token_num=int(seq_len.sum().item()),
+ max_q_seq_len=1,
+ max_kv_seq_len=max_kv_seq_len,
+ input_ids=input_ids.to(torch.int64).cpu(),
+ b_req_idx=req_idx,
+ b_mtp_index=mtp_index,
+ b_seq_len=seq_len,
+ mem_indexes_cpu=mem_indexes,
+ is_prefill=False,
+ multimodal_params=empty_multimodal_params(batch_size),
+ )
+
+ def _forward_prefill_input(self, model_input: ModelInput, allow_overlap: bool) -> ModelOutput:
+ if allow_overlap and self.args.enable_prefill_microbatch_overlap and model_input.batch_size > 1:
+ micro_input0, micro_input1 = self._split_prefill_input(model_input)
+ output0, output1 = self.model.microbatch_overlap_prefill(micro_input0, micro_input1)
+ return self._merge_model_outputs(output0, output1)
+ return self.model.forward(model_input)
+
+ def _forward_decode_input(self, model_input: ModelInput, allow_overlap: bool) -> ModelOutput:
+ if allow_overlap and self.args.enable_decode_microbatch_overlap and model_input.batch_size > 1:
+ micro_input0, micro_input1 = self._split_decode_input(model_input)
+ output0, output1 = self.model.microbatch_overlap_decode(micro_input0, micro_input1)
+ return self._merge_model_outputs(output0, output1)
+ return self.model.forward(model_input)
+
+ def _split_prefill_input(self, model_input: ModelInput):
+ split_batch = model_input.batch_size // 2
+ q_lens = model_input.b_seq_len - model_input.b_ready_cache_len
+ split_tokens = int(q_lens[:split_batch].sum().item())
+ return (
+ self._slice_prefill_input(model_input, 0, split_batch, 0, split_tokens),
+ self._slice_prefill_input(
+ model_input,
+ split_batch,
+ model_input.batch_size,
+ split_tokens,
+ int(q_lens.sum().item()),
+ ),
+ )
+
+ def _slice_prefill_input(
+ self,
+ model_input: ModelInput,
+ batch_start: int,
+ batch_end: int,
+ token_start: int,
+ token_end: int,
+ ) -> ModelInput:
+ b_seq_len = model_input.b_seq_len[batch_start:batch_end].clone()
+ b_ready_cache_len = model_input.b_ready_cache_len[batch_start:batch_end].clone()
+ b_q_seq_len = b_seq_len - b_ready_cache_len
+ b_prefill_start_loc = b_q_seq_len.cumsum(dim=0, dtype=torch.int32) - b_q_seq_len
+ has_output = model_input.b_prefill_has_output_cpu
+ return ModelInput(
+ batch_size=batch_end - batch_start,
+ total_token_num=int(b_seq_len.sum().item()),
+ max_q_seq_len=int(b_q_seq_len.max().item()),
+ max_kv_seq_len=int(b_seq_len.max().item()),
+ max_cache_len=int(b_ready_cache_len.max().item()),
+ prefix_total_token_num=int(b_ready_cache_len.sum().item()),
+ input_ids=model_input.input_ids[token_start:token_end].contiguous(),
+ b_req_idx=model_input.b_req_idx[batch_start:batch_end].clone(),
+ b_mtp_index=model_input.b_mtp_index[batch_start:batch_end].clone(),
+ b_seq_len=b_seq_len,
+ mem_indexes_cpu=model_input.mem_indexes_cpu[token_start:token_end].contiguous(),
+ is_prefill=True,
+ b_ready_cache_len=b_ready_cache_len,
+ b_prefill_start_loc=b_prefill_start_loc,
+ b_prefill_has_output_cpu=(has_output[batch_start:batch_end] if has_output is not None else None),
+ multimodal_params=model_input.multimodal_params[batch_start:batch_end],
+ )
+
+ def _split_decode_input(self, model_input: ModelInput):
+ split_batch = model_input.batch_size // 2
+ return (
+ self._slice_decode_input(model_input, 0, split_batch),
+ self._slice_decode_input(model_input, split_batch, model_input.batch_size),
+ )
+
+ def _slice_decode_input(self, model_input: ModelInput, batch_start: int, batch_end: int) -> ModelInput:
+ b_seq_len = model_input.b_seq_len[batch_start:batch_end].clone()
+ input_ids = model_input.input_ids
+ if input_ids is not None:
+ input_ids = input_ids[batch_start:batch_end].contiguous()
+ return ModelInput(
+ batch_size=batch_end - batch_start,
+ total_token_num=int(b_seq_len.sum().item()),
+ max_q_seq_len=model_input.max_q_seq_len,
+ max_kv_seq_len=int(b_seq_len.max().item()),
+ input_ids=input_ids,
+ b_req_idx=model_input.b_req_idx[batch_start:batch_end].clone(),
+ b_mtp_index=model_input.b_mtp_index[batch_start:batch_end].clone(),
+ b_seq_len=b_seq_len,
+ mem_indexes_cpu=model_input.mem_indexes_cpu[batch_start:batch_end].contiguous(),
+ is_prefill=False,
+ multimodal_params=model_input.multimodal_params[batch_start:batch_end],
+ )
+
+ def _merge_model_outputs(self, output0: ModelOutput, output1: ModelOutput) -> ModelOutput:
+ mtp_hiddens = None
+ if output0.mtp_main_output_hiddens is not None and output1.mtp_main_output_hiddens is not None:
+ mtp_hiddens = torch.cat(
+ (output0.mtp_main_output_hiddens, output1.mtp_main_output_hiddens),
+ dim=0,
+ )
+ return ModelOutput(
+ logits=torch.cat((output0.logits, output1.logits), dim=0),
+ prefill_mem_indexes_ready_event=output0.prefill_mem_indexes_ready_event,
+ mtp_main_output_hiddens=mtp_hiddens,
+ )
+
+ def _build_mtp_decode_index_tensors(self, req_idx: torch.Tensor, step_width: int):
+ batch_size = int(req_idx.shape[0])
+ return (
+ req_idx.repeat_interleave(step_width).to(torch.int32).cpu(),
+ torch.arange(step_width, dtype=torch.int32).repeat(batch_size),
+ )
+
+ def _build_mtp_seq_len(self, base_seq_len: torch.Tensor, step_width: int) -> torch.Tensor:
+ offsets = torch.arange(1, step_width + 1, dtype=torch.int32)
+ return (base_seq_len[:, None].to(torch.int32) + offsets[None, :]).reshape(-1)
+
+ def _alloc_req_indexes(self, batch_size: int) -> torch.Tensor:
+ req_indexes = [self.model.req_manager.alloc() for _ in range(batch_size)]
+ if any(index is None for index in req_indexes):
+ raise RuntimeError(f"failed to allocate {batch_size} request indexes")
+ return torch.tensor(req_indexes, dtype=torch.int32, device="cpu")
+
+ def _reset_model_cache(self):
+ self.model.mem_manager.free_all()
+ self.model.req_manager.free_all()
+ torch.cuda.synchronize()
+ torch.cuda.empty_cache()
+
+ def _argmax_ids(self, logits: torch.Tensor) -> torch.Tensor:
+ return torch.argmax(logits, dim=-1).detach().cpu().to(torch.int64)
+
+ def _touch_output(self, output: Optional[ModelOutput]):
+ if output is not None and output.logits is not None:
+ _ = output.logits.shape
+
+ def _make_result(
+ self,
+ case: BenchmarkCase,
+ elapsed_s: float,
+ measured_tokens: int,
+ warmup: bool,
+ ttft_elapsed_s: Optional[float] = None,
+ inter_token_latency_ms: Optional[float] = None,
+ ) -> BenchmarkResult:
+ """Convert raw timings into reported TPS and latency metrics."""
+ iters = self._case_iters(warmup)
+ scaled_tokens = measured_tokens * iters
+ qps = case.batch_size * iters / elapsed_s if elapsed_s > 0 else 0.0
+ tps = scaled_tokens / elapsed_s if elapsed_s > 0 else 0.0
+ ttft_ms = ttft_elapsed_s * 1000.0 / max(1, iters) if ttft_elapsed_s is not None else None
+ logical_tps = None
+ prefill_uncached_len = case.prefill_uncached_len
+ prefill_cached_len = None
+ if case.stage == "prefill":
+ uncached_len = int(case.prefill_uncached_len or case.context_len)
+ prefill_uncached_len = uncached_len
+ prefill_cached_len = max(0, case.context_len - uncached_len)
+ token_count = case.batch_size * case.context_len * iters
+ logical_tps = token_count / elapsed_s if elapsed_s > 0 else 0.0
+ return BenchmarkResult(
+ case=case.name,
+ stage=case.stage,
+ batch_size=case.batch_size,
+ context_len=case.context_len,
+ output_len=case.output_len,
+ chunked_prefill_size=case.chunked_prefill_size,
+ elapsed_ms=elapsed_s * 1000.0,
+ measured_tokens=scaled_tokens,
+ qps=qps,
+ tps=tps,
+ profiled_max_total_token_num=case.profiled_max_total_token_num,
+ profiled_batch_divisor=case.profiled_batch_divisor,
+ ttft_ms=ttft_ms,
+ inter_token_latency_ms=inter_token_latency_ms,
+ cache_hit_rate=case.cache_hit_rate,
+ prefill_uncached_len=prefill_uncached_len,
+ prefill_cached_len=prefill_cached_len,
+ prefill_step_tokens_per_req=case.prefill_step_tokens_per_req,
+ mtp_accept_rate=(
+ float(self.args.mtp_accept_rate) if case.stage == "decode" and self._mtp_enabled() else None
+ ),
+ logical_tps=logical_tps,
+ )
+
+ def _mtp_enabled(self) -> bool:
+ return self.args.mtp_mode in MTP_MODES and self.args.mtp_step > 0
+
+ def _mtp_step_width(self) -> int:
+ return int(self.args.mtp_step) + 1
+
+ def _num_mtp_modules(self) -> int:
+ if not self._mtp_enabled():
+ return 0
+ if self.args.mtp_mode.startswith("eagle"):
+ return 1
+ return int(self.args.mtp_step)
+
+
+def parse_typed_list(value: Optional[str], fallback: Sequence, cast) -> List:
+ if value is None or value == "":
+ return list(fallback)
+ if isinstance(value, cast):
+ return [value]
+ normalized = str(value).replace(",", " ")
+ return [cast(item) for item in normalized.split() if item.strip()]
+
+
+def parse_int_list(value: Optional[str], fallback: Sequence[int]) -> List[int]:
+ return parse_typed_list(value, fallback, int)
+
+
+def parse_float_list(value: Optional[str], fallback: Sequence[float]) -> List[float]:
+ return parse_typed_list(value, fallback, float)
+
+
+def parse_chunk_sizes(value: Optional[str], fallback: Optional[int]) -> List[Optional[int]]:
+ if value is None:
+ return [fallback] if fallback else [None]
+ chunks: List[Optional[int]] = []
+ for item in str(value).replace(",", " ").split():
+ item = item.strip().lower()
+ if item in {"none", "full", "0", "-1"}:
+ chunks.append(None)
+ else:
+ chunks.append(int(item))
+ return chunks or [None]
+
+
+def prefill_uncached_len(context_len: int, cache_hit_rate: float) -> int:
+ """Return the uncached suffix length for a prompt-cache hit ratio."""
+ if cache_hit_rate < 0.0 or cache_hit_rate >= 1.0:
+ raise ValueError(f"cache hit rate must satisfy 0 <= hit < 1, got {cache_hit_rate}")
+ uncached = int(math.ceil(context_len * (1.0 - cache_hit_rate)))
+ return max(1, min(context_len, uncached))
+
+
+def prefill_step_tokens_per_req(uncached_len: int, chunked_prefill_size: Optional[int]) -> int:
+ """Return tokens handled per request in one production prefill step."""
+ if chunked_prefill_size and chunked_prefill_size > 0:
+ return max(1, min(uncached_len, int(chunked_prefill_size)))
+ return max(1, uncached_len)
+
+
+def format_cache_hit_suffix(cache_hit_rate: float) -> str:
+ return f"{cache_hit_rate:.4f}".rstrip("0").rstrip(".").replace(".", "p")
+
+
+def apply_max_batch_size(batch_size: int, max_batch_size: int) -> int:
+ """Apply the benchmark-wide auto batch-size upper bound."""
+ if max_batch_size > 0:
+ batch_size = min(batch_size, int(max_batch_size))
+ return max(1, batch_size)
+
+
+def prefill_batch_size_from_batch_max_tokens(
+ batch_max_tokens: int,
+ step_tokens_per_req: int,
+ max_batch_size: int,
+) -> int:
+ """Compute prefill BS from batch_max_tokens before KV-capacity capping."""
+ batch_size = max(1, int(batch_max_tokens) // max(1, step_tokens_per_req))
+ return apply_max_batch_size(batch_size, max_batch_size)
+
+
+def build_prefill_cases(
+ args: SimpleNamespace,
+ input_lens: Sequence[int],
+ chunk_sizes: Sequence[Optional[int]],
+ cache_hit_rates: Sequence[float],
+) -> List[BenchmarkCase]:
+ """Build full-prefill cases using batch_max_tokens per chunk step."""
+ if args.batch_max_tokens is None:
+ raise ValueError("prefill benchmark requires --batch_max_tokens")
+ cases: List[BenchmarkCase] = []
+ for input_len in input_lens:
+ for chunk_size in chunk_sizes:
+ for cache_hit_rate in cache_hit_rates:
+ uncached_len = prefill_uncached_len(input_len, cache_hit_rate)
+ step_tokens = prefill_step_tokens_per_req(uncached_len, chunk_size)
+ bs = prefill_batch_size_from_batch_max_tokens(
+ args.batch_max_tokens,
+ step_tokens,
+ args.max_batch_size,
+ )
+ chunk_name = chunk_size if chunk_size else "none"
+ hit_name = format_cache_hit_suffix(cache_hit_rate)
+ cases.append(
+ BenchmarkCase(
+ name=(
+ f"prefill_bs{bs}_in{input_len}_hit{hit_name}"
+ f"_uncached{uncached_len}_chunk{chunk_name}"
+ f"_btok{args.batch_max_tokens}"
+ ),
+ stage="prefill",
+ batch_size=bs,
+ context_len=input_len,
+ output_len=0,
+ chunked_prefill_size=chunk_size,
+ cache_hit_rate=cache_hit_rate,
+ prefill_uncached_len=uncached_len,
+ prefill_step_tokens_per_req=step_tokens,
+ prefill_batch_size_by_batch_max_tokens=bs,
+ )
+ )
+ return cases
+
+
+def build_decode_cases(
+ args: SimpleNamespace,
+ batch_sizes: Sequence[int],
+ context_lens: Sequence[int],
+ output_lens: Sequence[int],
+) -> List[BenchmarkCase]:
+ """Build decode cases; profile mode resolves BS after model load."""
+ decode_batch_sizes = [1] if args.decode_batch_size_mode == "profile" else batch_sizes
+ cases: List[BenchmarkCase] = []
+ for bs in decode_batch_sizes:
+ for context_len in context_lens:
+ for output_len in output_lens:
+ profile_suffix = "_profilebs" if args.decode_batch_size_mode == "profile" else ""
+ cases.append(
+ BenchmarkCase(
+ name=(f"decode_bs{bs}_ctx{context_len}_out{output_len}" f"{profile_suffix}"),
+ stage="decode",
+ batch_size=bs,
+ context_len=context_len,
+ output_len=output_len,
+ )
+ )
+ return cases
+
+
+def build_cases(args: SimpleNamespace) -> List[BenchmarkCase]:
+ """Expand CLI list options into concrete prefill/decode benchmark cases."""
+ batch_sizes = parse_int_list(args.batch_sizes, [args.batch_size] if args.batch_size else DEFAULT_BATCH_SIZES)
+ input_lens = parse_int_list(args.input_lens, [args.input_len])
+ context_lens = parse_int_list(args.context_lens, input_lens)
+ output_lens = parse_int_list(args.output_lens, [args.output_len])
+ chunk_sizes = parse_chunk_sizes(args.chunked_prefill_sizes, args.chunked_prefill_size)
+ cache_hit_rates = parse_float_list(args.prefill_cache_hit_rates, [0.0])
+
+ cases: List[BenchmarkCase] = []
+ if args.benchmark in {"all", "prefill"}:
+ cases.extend(build_prefill_cases(args, input_lens, chunk_sizes, cache_hit_rates))
+ if args.benchmark in {"all", "decode"}:
+ cases.extend(build_decode_cases(args, batch_sizes, context_lens, output_lens))
+ return cases
+
+
+def decode_profile_batch_divisor(args: SimpleNamespace, case: BenchmarkCase) -> int:
+ """Reserve KV capacity for context, generated tokens, and MTP expansion."""
+ mtp_width = int(args.mtp_step) + 1 if args.mtp_mode in MTP_MODES else 1
+ return max(1, case.context_len + case.output_len + mtp_width + 8)
+
+
+def filter_capacity_decode_cases(
+ args: SimpleNamespace,
+ cases: Sequence[BenchmarkCase],
+ profiled_max_total_token_num: int,
+) -> List[BenchmarkCase]:
+ if not getattr(args, "decode_filter_capacity", False):
+ return list(cases)
+
+ resolved: List[BenchmarkCase] = []
+ capacity_tokens = int(profiled_max_total_token_num)
+ for case in cases:
+ if case.stage != "decode":
+ resolved.append(case)
+ continue
+ divisor = decode_profile_batch_divisor(args, case)
+ if case.batch_size * divisor <= capacity_tokens:
+ resolved.append(
+ replace(
+ case,
+ profiled_max_total_token_num=capacity_tokens,
+ profiled_batch_divisor=divisor,
+ )
+ )
+ return resolved
+
+
+def resolve_profile_decode_cases(
+ args: SimpleNamespace,
+ cases: Sequence[BenchmarkCase],
+ profiled_max_total_token_num: int,
+) -> List[BenchmarkCase]:
+ """Replace decode profile placeholders with capacity-derived max BS."""
+ if args.decode_batch_size_mode != "profile":
+ return list(cases)
+
+ resolved: List[BenchmarkCase] = []
+ for case in cases:
+ if case.stage != "decode":
+ resolved.append(case)
+ continue
+
+ divisor = decode_profile_batch_divisor(args, case)
+ batch_size = max(1, int(profiled_max_total_token_num) // divisor)
+ batch_size = apply_max_batch_size(batch_size, args.max_batch_size)
+
+ resolved.append(
+ replace(
+ case,
+ name=(
+ f"decode_bs{batch_size}_ctx{case.context_len}"
+ f"_out{case.output_len}_profile{profiled_max_total_token_num}"
+ ),
+ batch_size=batch_size,
+ profiled_max_total_token_num=int(profiled_max_total_token_num),
+ profiled_batch_divisor=divisor,
+ )
+ )
+
+ return resolved
+
+
+def resolve_batch_max_prefill_cases(
+ args: SimpleNamespace,
+ cases: Sequence[BenchmarkCase],
+ profiled_max_total_token_num: int,
+) -> List[BenchmarkCase]:
+ """Cap prefill BS by profiled KV capacity after the model is loaded."""
+ resolved: List[BenchmarkCase] = []
+ capacity_tokens = int(profiled_max_total_token_num)
+ for case in cases:
+ if case.stage != "prefill":
+ resolved.append(case)
+ continue
+
+ if case.context_len <= 0:
+ raise ValueError(f"invalid prefill context_len={case.context_len}")
+ bs_by_capacity = capacity_tokens // case.context_len
+ if bs_by_capacity <= 0:
+ raise ValueError(
+ "single prefill request does not fit profiled token capacity: "
+ f"context_len={case.context_len} capacity={profiled_max_total_token_num}"
+ )
+
+ bs_by_batch = int(case.prefill_batch_size_by_batch_max_tokens or case.batch_size)
+ batch_size = min(bs_by_batch, bs_by_capacity)
+ batch_size = apply_max_batch_size(batch_size, args.max_batch_size)
+
+ chunk_name = case.chunked_prefill_size if case.chunked_prefill_size else "none"
+ hit_name = format_cache_hit_suffix(case.cache_hit_rate)
+ resolved.append(
+ replace(
+ case,
+ name=(
+ f"prefill_bs{batch_size}_in{case.context_len}_hit{hit_name}"
+ f"_uncached{case.prefill_uncached_len}_chunk{chunk_name}"
+ f"_btok{args.batch_max_tokens}"
+ f"_cap{profiled_max_total_token_num}"
+ ),
+ batch_size=batch_size,
+ profiled_max_total_token_num=int(profiled_max_total_token_num),
+ profiled_batch_divisor=case.context_len,
+ )
+ )
+
+ return resolved
+
+
+def normalize_args(args: argparse.Namespace, cases: Sequence[BenchmarkCase]) -> SimpleNamespace:
+ """Fill LightLLM startup args needed before model construction."""
+ if args.data_type is None:
+ args.data_type = get_dtype(args.model_dir)
+
+ if args.quant_type is None:
+ args.quant_type = "none"
+
+ if not 0.0 <= float(args.mtp_accept_rate) <= 1.0:
+ raise ValueError(f"--mtp_accept_rate must be in [0, 1], got {args.mtp_accept_rate}")
+
+ max_batch = max(case.batch_size for case in cases)
+ max_context = max(case.context_len for case in cases)
+ max_output = max(case.output_len for case in cases)
+ mtp_width = (args.mtp_step + 1) if args.mtp_mode in MTP_MODES else 1
+ max_runtime_len = max_context + max_output + mtp_width + 2
+
+ if args.max_req_total_len is None:
+ args.max_req_total_len = max_runtime_len
+ else:
+ args.max_req_total_len = max(args.max_req_total_len, max_runtime_len)
+
+ if args.graph_max_len_in_batch == 0:
+ args.graph_max_len_in_batch = args.max_req_total_len
+
+ max_prefill_chunk = (
+ max(
+ min(case.context_len, case.chunked_prefill_size or case.context_len)
+ for case in cases
+ if case.stage == "prefill"
+ )
+ if any(case.stage == "prefill" for case in cases)
+ else max_context
+ )
+ if args.batch_max_tokens is None:
+ args.batch_max_tokens = max(max_batch * max_prefill_chunk, max_batch * mtp_width, 1)
+
+ decode_batch_size_needs_profile = (
+ args.benchmark in {"all", "decode"}
+ and args.decode_batch_size_mode == "profile"
+ and args.max_total_token_num is None
+ )
+ prefill_batch_size_needs_profile = args.benchmark in {"all", "prefill"} and args.max_total_token_num is None
+ needs_profiled_batch_size = decode_batch_size_needs_profile or prefill_batch_size_needs_profile
+
+ if args.max_total_token_num is None and not needs_profiled_batch_size:
+ args.max_total_token_num = max_batch * (args.max_req_total_len + mtp_width + 8)
+ if args.max_total_token_num is not None:
+ args.max_total_token_num = max(args.max_total_token_num, args.batch_max_tokens + 1, args.max_req_total_len)
+
+ if decode_batch_size_needs_profile and args.max_batch_size > 0:
+ args.running_max_req_size = max(args.running_max_req_size, int(args.max_batch_size))
+ # Profile decode BS is resolved after model load. Use the cap as the
+ # pre-load upper bound so request slots and optional decode graphs agree.
+ if not args.disable_cudagraph:
+ args.graph_max_batch_size = max(args.graph_max_batch_size, int(args.max_batch_size))
+ if prefill_batch_size_needs_profile:
+ args.running_max_req_size = max(args.running_max_req_size, max_batch)
+
+ if args.graph_max_batch_size < max_batch:
+ args.graph_max_batch_size = max_batch
+
+ if args.nccl_port is None:
+ args.nccl_port = 28765
+
+ if args.mtp_mode in MTP_MODES:
+ if args.mtp_step <= 0:
+ raise ValueError("--mtp_mode requires --mtp_step > 0")
+ if not args.mtp_draft_model_dir:
+ raise ValueError("--mtp_mode requires --mtp_draft_model_dir")
+ args.mtp_draft_model_dir = normalize_mtp_draft_dirs(args.mtp_mode, args.mtp_step, args.mtp_draft_model_dir)
+ else:
+ args.mtp_mode = None
+ args.mtp_step = 0
+ args.mtp_draft_model_dir = None
+
+ return SimpleNamespace(**vars(args))
+
+
+def normalize_mtp_draft_dirs(mtp_mode: str, mtp_step: int, draft_dirs: Sequence[str]) -> List[str]:
+ expected = 1 if mtp_mode.startswith("eagle") else mtp_step
+ if isinstance(draft_dirs, str):
+ draft_dirs = [draft_dirs]
+ draft_dirs = list(draft_dirs)
+ if len(draft_dirs) == 1 and expected > 1:
+ return draft_dirs * expected
+ if len(draft_dirs) != expected:
+ raise ValueError(f"{mtp_mode} expects {expected} draft model dir(s), got {len(draft_dirs)}")
+ return draft_dirs
+
+
+def build_model_kvargs(args: SimpleNamespace, rank_id: int) -> Dict:
+ return {
+ "args": args,
+ "nccl_host": args.nccl_host,
+ "nccl_port": args.nccl_port,
+ "rank_id": rank_id,
+ "world_size": args.tp,
+ "dp_size": args.dp,
+ "weight_dir": args.model_dir,
+ "data_type": args.data_type,
+ "quant_type": args.quant_type,
+ "quant_cfg": args.quant_cfg,
+ "expert_dtype": args.expert_dtype,
+ "load_way": "HF",
+ "max_total_token_num": args.max_total_token_num,
+ "graph_max_len_in_batch": args.graph_max_len_in_batch,
+ "graph_max_batch_size": args.graph_max_batch_size,
+ "mem_fraction": args.mem_fraction,
+ "max_req_num": max(args.running_max_req_size, args.graph_max_batch_size),
+ "batch_max_tokens": args.batch_max_tokens,
+ "run_mode": "normal",
+ "max_seq_length": args.max_req_total_len,
+ "disable_cudagraph": args.disable_cudagraph,
+ "llm_prefill_att_backend": args.llm_prefill_att_backend,
+ "llm_decode_att_backend": args.llm_decode_att_backend,
+ "vit_att_backend": args.vit_att_backend,
+ "llm_kv_type": args.llm_kv_type,
+ "llm_kv_quant_group_size": args.llm_kv_quant_group_size,
+ }
+
+
+def init_mtp_draft_models(args: SimpleNamespace, main_kvargs: Dict, main_model) -> List:
+ if args.mtp_mode not in MTP_MODES:
+ return []
+
+ os.environ["DISABLE_CHECK_MAX_LEN_INFER"] = "1"
+ draft_models = []
+ for draft_dir in args.mtp_draft_model_dir:
+ mtp_cfg, _ = PretrainedConfig.get_config_dict(draft_dir)
+ model_type = mtp_cfg.get("model_type", "")
+ mtp_kvargs = {
+ "weight_dir": draft_dir,
+ "max_total_token_num": main_model.mem_manager.size,
+ "load_way": main_kvargs["load_way"],
+ "max_req_num": main_kvargs["max_req_num"],
+ "max_seq_length": main_kvargs["max_seq_length"],
+ "is_token_healing": False,
+ "return_all_prompt_logics": False,
+ "disable_chunked_prefill": args.disable_chunked_prefill,
+ "data_type": main_kvargs["data_type"],
+ "graph_max_batch_size": main_kvargs["graph_max_batch_size"],
+ "graph_max_len_in_batch": main_kvargs["graph_max_len_in_batch"],
+ "disable_cudagraph": main_kvargs["disable_cudagraph"],
+ "mem_fraction": main_kvargs["mem_fraction"],
+ "batch_max_tokens": main_kvargs["batch_max_tokens"],
+ "quant_type": main_kvargs["quant_type"],
+ "quant_cfg": main_kvargs["quant_cfg"],
+ "expert_dtype": main_kvargs["expert_dtype"],
+ "run_mode": "normal",
+ "main_model": main_model,
+ "mtp_previous_draft_models": draft_models.copy(),
+ }
+ if model_type == "deepseek_v3":
+ assert args.mtp_mode in {
+ "vanilla_with_att",
+ "eagle_with_att",
+ }, f"{model_type} MTP requires *_with_att mode"
+ draft_models.append(Deepseek3MTPModel(mtp_kvargs))
+ elif model_type == "deepseek_v4":
+ assert args.mtp_mode == "eagle_with_att", f"{model_type} MTP requires eagle_with_att mode"
+ draft_models.append(DeepseekV4MTPModel(mtp_kvargs))
+ elif model_type == "qwen3_moe":
+ assert args.mtp_mode in {
+ "vanilla_no_att",
+ "eagle_no_att",
+ }, f"{model_type} MTP requires *_no_att mode"
+ draft_models.append(Qwen3MOEMTPModel(mtp_kvargs))
+ elif model_type == "mistral":
+ assert args.mtp_mode in {
+ "vanilla_no_att",
+ "eagle_no_att",
+ }, f"{model_type} MTP requires *_no_att mode"
+ draft_models.append(MistralMTPModel(mtp_kvargs))
+ elif model_type == "glm4_moe_lite":
+ assert args.mtp_mode in {
+ "vanilla_with_att",
+ "eagle_with_att",
+ }, f"{model_type} MTP requires *_with_att mode"
+ draft_models.append(Glm4MoeLiteMTPModel(mtp_kvargs))
+ else:
+ raise ValueError(f"unsupported MTP draft model_type={model_type} from {draft_dir}")
+ return draft_models
+
+
+def run_worker(args_dict: Dict, case_dicts: List[Dict], rank_id: int, ans_queue):
+ try:
+ args = SimpleNamespace(**args_dict)
+ cases = [BenchmarkCase(**case) for case in case_dicts]
+ set_env_start_args(args)
+
+ from lightllm.distributed import dist_group_manager
+ import torch.distributed as dist
+
+ model_kvargs = build_model_kvargs(args, rank_id)
+ group_size = 2 if (args.enable_decode_microbatch_overlap or args.enable_prefill_microbatch_overlap) else 1
+ if group_size == 2:
+ for case in cases:
+ assert case.batch_size % 2 == 0, "microbatch overlap requires even batch_size"
+
+ init_distributed_env(model_kvargs)
+ dist_group_manager.create_groups(group_size=group_size)
+ model_cfg, _ = PretrainedConfig.get_config_dict(args.model_dir)
+ dist.barrier()
+
+ torch.cuda.empty_cache()
+ model, _ = get_model(model_cfg, model_kvargs)
+ cases = resolve_batch_max_prefill_cases(args, cases, model.mem_manager.size)
+ cases = resolve_profile_decode_cases(args, cases, model.mem_manager.size)
+ cases = filter_capacity_decode_cases(args, cases, model.mem_manager.size)
+ if not cases:
+ raise ValueError("no benchmark cases remain after capacity filtering")
+ draft_models = init_mtp_draft_models(args, model_kvargs, model)
+ token_source = TokenSource(args)
+ executor = StaticBenchmarkExecutor(args, model, draft_models, token_source)
+
+ results = []
+ log_progress = rank_id == args.node_rank * args.tp
+ for case_index, case in enumerate(cases, start=1):
+ if log_progress:
+ print(f"[rank {rank_id}] case {case_index}/{len(cases)} start {case.name}", flush=True)
+ if args.warmup_iters > 0:
+ executor.run_case(case, warmup=True)
+ result = executor.run_case(case, warmup=False)
+ results.append(asdict(result))
+ if log_progress:
+ itl = "" if result.inter_token_latency_ms is None else f" itl_ms={result.inter_token_latency_ms:.3f}"
+ print(
+ f"[rank {rank_id}] case {case_index}/{len(cases)} done elapsed_ms={result.elapsed_ms:.3f}{itl}",
+ flush=True,
+ )
+ dist.barrier()
+
+ ans_queue.put({"ok": True, "rank": rank_id, "results": results})
+ except Exception:
+ ans_queue.put({"ok": False, "rank": rank_id, "traceback": traceback.format_exc()})
+ finally:
+ try:
+ ans_queue.close()
+ ans_queue.join_thread()
+ except Exception:
+ pass
+ os._exit(0)
+
+
+def fmt_optional(value, precision: int = 2) -> str:
+ if value is None:
+ return "-"
+ if isinstance(value, float):
+ return f"{value:.{precision}f}"
+ return str(value)
+
+
+def print_aligned_table(headers: Sequence[str], rows: Sequence[Sequence[str]]):
+ """Print a compact right-aligned ASCII table."""
+ if not rows:
+ return
+ widths = [len(str(header)) for header in headers]
+ for row in rows:
+ for index, value in enumerate(row):
+ widths[index] = max(widths[index], len(str(value)))
+
+ def format_row(row: Sequence[str]) -> str:
+ return " ".join(str(value).rjust(widths[index]) for index, value in enumerate(row))
+
+ print(format_row(headers), flush=True)
+ print(" ".join("-" * width for width in widths), flush=True)
+ for row in rows:
+ print(format_row(row), flush=True)
+
+
+def prefill_table_row(result: BenchmarkResult) -> List[str]:
+ """Format one prefill result row for stdout table output."""
+ return [
+ str(result.context_len),
+ f"{result.cache_hit_rate:.2f}",
+ str(result.batch_size),
+ fmt_optional(result.profiled_max_total_token_num, 0),
+ fmt_optional(result.prefill_uncached_len, 0),
+ fmt_optional(result.prefill_cached_len, 0),
+ str(result.measured_tokens),
+ f"{result.elapsed_ms:.3f}",
+ f"{result.qps:.2f}",
+ f"{result.tps:.2f}",
+ fmt_optional(result.logical_tps, 2),
+ ]
+
+
+def decode_table_row(result: BenchmarkResult) -> List[str]:
+ """Format one decode result row for stdout table output."""
+ return [
+ str(result.context_len),
+ str(result.batch_size),
+ fmt_optional(result.mtp_accept_rate, 2),
+ fmt_optional(result.profiled_max_total_token_num, 0),
+ f"{result.elapsed_ms:.3f}",
+ f"{result.qps:.2f}",
+ f"{result.tps:.2f}",
+ fmt_optional(result.inter_token_latency_ms, 3),
+ ]
+
+
+def print_results_table(results: Sequence[BenchmarkResult]):
+ """Print separate prefill/decode tables for measured results."""
+ prefill_rows = [prefill_table_row(result) for result in results if result.stage == "prefill"]
+ decode_rows = [decode_table_row(result) for result in results if result.stage == "decode"]
+
+ if prefill_rows:
+ print("\n[prefill]", flush=True)
+ print_aligned_table(PREFILL_TABLE_HEADERS, prefill_rows)
+ if decode_rows:
+ print("\n[decode]", flush=True)
+ print_aligned_table(DECODE_TABLE_HEADERS, decode_rows)
+
+
+def _dp_size(args: SimpleNamespace) -> int:
+ return max(1, int(args.dp or 1))
+
+
+def _dp_world_size(args: SimpleNamespace) -> int:
+ return max(1, int(args.tp) // _dp_size(args))
+
+
+def _is_dp_group_leader(args: SimpleNamespace, rank_id: int) -> bool:
+ return rank_id % _dp_world_size(args) == 0
+
+
+def _raw_decode_step_count(result: Dict) -> int:
+ inter_token_latency_ms = result.get("inter_token_latency_ms")
+ if inter_token_latency_ms is None or inter_token_latency_ms <= 0:
+ return 0
+ return max(1, int(round(float(result["elapsed_ms"]) / float(inter_token_latency_ms))))
+
+
+def aggregate_rank_results(args: SimpleNamespace, messages: Sequence[Dict]) -> List[Dict]:
+ """Aggregate rank-local measurements into one global result per case."""
+ by_case: Dict[int, List[Dict]] = {}
+ for message in messages:
+ rank_id = int(message["rank"])
+ for case_index, result in enumerate(message.get("results") or []):
+ by_case.setdefault(case_index, []).append({"rank": rank_id, "result": result})
+
+ aggregated_results: List[Dict] = []
+ iters = int(args.bench_iters)
+ for case_index in sorted(by_case):
+ rank_items = sorted(by_case[case_index], key=lambda item: int(item["rank"]))
+ leader_items = [item for item in rank_items if _is_dp_group_leader(args, int(item["rank"]))]
+ if not leader_items:
+ leader_items = rank_items
+
+ first = dict(leader_items[0]["result"])
+ elapsed_ms = max(float(item["result"]["elapsed_ms"]) for item in rank_items)
+ elapsed_s = elapsed_ms / 1000.0
+ batch_size = sum(int(item["result"]["batch_size"]) for item in leader_items)
+ measured_tokens = sum(int(item["result"]["measured_tokens"]) for item in leader_items)
+
+ first["batch_size"] = batch_size
+ first["elapsed_ms"] = elapsed_ms
+ first["measured_tokens"] = measured_tokens
+ first["qps"] = batch_size * iters / elapsed_s if elapsed_s > 0 else 0.0
+ first["tps"] = measured_tokens / elapsed_s if elapsed_s > 0 else 0.0
+
+ profiled_values = [
+ int(item["result"]["profiled_max_total_token_num"])
+ for item in leader_items
+ if item["result"].get("profiled_max_total_token_num") is not None
+ ]
+ if profiled_values:
+ first["profiled_max_total_token_num"] = sum(profiled_values)
+
+ if first["stage"] == "prefill":
+ logical_tokens = batch_size * int(first["context_len"]) * iters
+ first["logical_tps"] = logical_tokens / elapsed_s if elapsed_s > 0 else 0.0
+
+ slowest_item = max(rank_items, key=lambda item: float(item["result"]["elapsed_ms"]))
+ decode_step_count = _raw_decode_step_count(slowest_item["result"])
+ if decode_step_count > 0:
+ first["inter_token_latency_ms"] = elapsed_ms / decode_step_count
+
+ ttft_values = [
+ float(item["result"]["ttft_ms"]) for item in rank_items if item["result"].get("ttft_ms") is not None
+ ]
+ if ttft_values:
+ first["ttft_ms"] = max(ttft_values)
+
+ aggregated_results.append(first)
+
+ return aggregated_results
+
+
+def run_benchmark(args: SimpleNamespace, cases: Sequence[BenchmarkCase]) -> List[Dict]:
+ ctx = mp.get_context("spawn")
+ ans_queue = ctx.Queue()
+ workers = []
+ rank_start = args.node_rank * args.tp
+ rank_end = (args.node_rank + 1) * args.tp
+ case_dicts = [asdict(case) for case in cases]
+ args_dict = vars(args)
+
+ for rank_id in range(rank_start, rank_end):
+ proc = ctx.Process(target=run_worker, args=(args_dict, case_dicts, rank_id, ans_queue))
+ proc.start()
+ workers.append(proc)
+
+ messages = []
+ while len(messages) < len(workers):
+ try:
+ messages.append(ans_queue.get(timeout=5))
+ continue
+ except queue.Empty:
+ if all(not proc.is_alive() for proc in workers):
+ break
+
+ for proc in workers:
+ proc.join()
+
+ failed = [message for message in messages if not message.get("ok")]
+ reported_ranks = {int(message["rank"]) for message in messages if "rank" in message}
+ failed.extend(
+ {
+ "ok": False,
+ "rank": rank_start + index,
+ "traceback": f"worker exited with code {proc.exitcode}",
+ }
+ for index, proc in enumerate(workers)
+ if proc.exitcode not in (0, None)
+ )
+ failed.extend(
+ {
+ "ok": False,
+ "rank": rank_start + index,
+ "traceback": "worker did not report a result",
+ }
+ for index, proc in enumerate(workers)
+ if rank_start + index not in reported_ranks and proc.exitcode in (0, None)
+ )
+ if failed:
+ for item in failed:
+ print(
+ f"rank {item.get('rank')} failed:\n{item.get('traceback')}",
+ file=sys.stderr,
+ )
+ raise RuntimeError(f"{len(failed)} worker(s) failed")
+
+ results = aggregate_rank_results(args, messages)
+ result_objs = [BenchmarkResult(**result) for result in results]
+ print_results_table(result_objs)
+ return results
+
+
+def add_static_benchmark_args(parser: argparse.ArgumentParser):
+ parser.add_argument("--benchmark", choices=["all", "prefill", "decode"], default="all")
+ parser.add_argument("--batch_size", type=int, default=None, help="legacy single batch size")
+ parser.add_argument(
+ "--batch_sizes",
+ type=str,
+ default=None,
+ help="comma/space separated batch sizes",
+ )
+ parser.add_argument("--input_len", type=int, default=64, help="legacy single prefill/context length")
+ parser.add_argument(
+ "--input_lens",
+ type=str,
+ default=None,
+ help="comma/space separated prefill lengths",
+ )
+ parser.add_argument(
+ "--context_lens",
+ type=str,
+ default=None,
+ help="comma/space separated decode context lengths",
+ )
+ parser.add_argument("--output_len", type=int, default=512, help="legacy single decode output length")
+ parser.add_argument(
+ "--output_lens",
+ type=str,
+ default=None,
+ help="comma/space separated decode output lengths",
+ )
+ parser.add_argument(
+ "--chunked_prefill_sizes",
+ type=str,
+ default=4096,
+ help=("comma/space separated prefill chunk sizes; default is 4096 " "(full/none/0 select unchunked prefill)"),
+ )
+ parser.add_argument(
+ "--prefill_cache_hit_rates",
+ type=str,
+ default=None,
+ help=(
+ "comma/space separated cache hit rates for prefill, e.g. "
+ "'0,0.5,0.8,0.9'; uncached tokens are ceil(input_len * (1-hit))"
+ ),
+ )
+ parser.add_argument(
+ "--max_batch_size",
+ type=int,
+ default=2048,
+ help="upper bound for auto-computed prefill/decode batch size; <=0 disables it",
+ )
+ parser.add_argument(
+ "--decode_batch_size_mode",
+ choices=["explicit", "profile"],
+ default="explicit",
+ help=(
+ "explicit uses --batch_size/--batch_sizes; profile computes decode BS "
+ "from profiled max_total_token_num per context"
+ ),
+ )
+ parser.add_argument(
+ "--decode_filter_capacity",
+ action="store_true",
+ help="drop explicit decode cases whose batch size cannot fit profiled KV capacity",
+ )
+ parser.add_argument(
+ "--mtp_accept_rate",
+ type=float,
+ default=1.0,
+ help=("per-draft-token MTP acceptance probability; sampling is outside " "the timed decode section"),
+ )
+ parser.add_argument("--warmup_iters", type=int, default=1)
+ parser.add_argument("--bench_iters", type=int, default=1)
+ parser.add_argument("--seed", type=int, default=1234)
+ parser.add_argument("--dump_file", type=str, default=None, help="write aggregated benchmark results as JSON")
+
+
+def main(argv: Optional[Sequence[str]] = None):
+ parser = make_argument_parser()
+ add_static_benchmark_args(parser)
+ args = parser.parse_args(argv)
+ if args.benchmark in {"all", "prefill"} and args.batch_max_tokens is None:
+ args.batch_max_tokens = 8192
+ cases = build_cases(args)
+ if not cases:
+ raise ValueError("no benchmark cases were generated")
+ args = normalize_args(args, cases)
+ set_env_start_args(args)
+
+ results = run_benchmark(args, cases)
+ if args.dump_file:
+ dump_path = Path(args.dump_file)
+ dump_path.parent.mkdir(parents=True, exist_ok=True)
+ payload = {"args": vars(args), "results": results}
+ dump_path.write_text(json.dumps(payload, indent=2, sort_keys=True))
+
+
+if __name__ == "__main__":
+ mp.set_start_method("spawn", force=True)
+ main()