diff --git a/docs/CN/source/cookbook/sensenova_loop_deployment.rst b/docs/CN/source/cookbook/sensenova_loop_deployment.rst new file mode 100644 index 0000000000..26bf5951f2 --- /dev/null +++ b/docs/CN/source/cookbook/sensenova_loop_deployment.rst @@ -0,0 +1,60 @@ +.. _sensenova_loop_deployment: + +SenseNova Loop 部署指南 +=========================== + +启动服务 +-------- + +下面以 4 卡部署为例,在仓库根目录执行: + +.. code-block:: bash + + CUDA_VISIBLE_DEVICES=0,1,2,3 LOADWORKER=4 LIGHTLLM_TRITON_AUTOTUNE_LEVEL=1 \ + python -m lightllm.server.api_server \ + --model_dir /path/to/SenseNova-Loop \ + --model_name sensenova_loop \ + --tp 4 \ + --port 8000 + +启用 MTP +-------- + +checkpoint 已包含 MTP 权重,增加以下参数即可启用: + +.. code-block:: bash + + CUDA_VISIBLE_DEVICES=0,1,2,3 LOADWORKER=4 LIGHTLLM_TRITON_AUTOTUNE_LEVEL=1 \ + python -m lightllm.server.api_server \ + --model_dir /path/to/SenseNova-Loop \ + --model_name sensenova_loop \ + --mtp_mode eagle_with_att \ + --mtp_step 3 \ + --tp 4 \ + --port 8000 + +2P1D 部署 +--------- + +下面的脚本默认使用 8 张 GPU,启动两个 TP2 Prefill 服务和一个 TP4 Decode 服务,并启用 MTP: + +.. code-block:: bash + + MODEL_DIR=/path/to/SenseNova-Loop PORT=8000 \ + bash test/start_scripts/sensenova_loop/sensenova_loop_pd_2p1d.sh + +可通过 ``GPU_IDS`` 指定显卡列表,添加 ``DRY_RUN=1`` 可先查看启动命令。 +按 ``Ctrl+C`` 会关闭所有服务。 + +请求示例 +-------- + +.. code-block:: bash + + curl http://localhost:8000/v1/chat/completions \ + -H 'Content-Type: application/json' \ + -d '{ + "model": "sensenova_loop", + "messages": [{"role": "user", "content": "你好,请用一句话介绍自己。"}], + "max_tokens": 128 + }' diff --git a/docs/CN/source/index.rst b/docs/CN/source/index.rst index 4c7fa3077c..a0399f1623 100755 --- a/docs/CN/source/index.rst +++ b/docs/CN/source/index.rst @@ -67,6 +67,7 @@ Lightllm 整合了众多的开源方案的优点,包括但不限于 FasterTran GLM-4.7-Flash 部署 Qwen3.5 部署 + SenseNova Loop 部署 .. toctree:: :maxdepth: 1 diff --git a/docs/EN/source/cookbook/sensenova_loop_deployment.rst b/docs/EN/source/cookbook/sensenova_loop_deployment.rst new file mode 100644 index 0000000000..4190403f8d --- /dev/null +++ b/docs/EN/source/cookbook/sensenova_loop_deployment.rst @@ -0,0 +1,60 @@ +.. _sensenova_loop_deployment: + +SenseNova Loop Deployment Guide +============================== + +Start the Server +---------------- + +This example starts a BF16 service on four GPUs. Run it from the repository root: + +.. code-block:: bash + + CUDA_VISIBLE_DEVICES=0,1,2,3 LOADWORKER=4 LIGHTLLM_TRITON_AUTOTUNE_LEVEL=1 \ + python -m lightllm.server.api_server \ + --model_dir /path/to/SenseNova-Loop \ + --model_name sensenova_loop \ + --tp 4 \ + --port 8000 + +Enable MTP +---------- + +The checkpoint includes MTP weights. Add the following options to enable MTP: + +.. code-block:: bash + + CUDA_VISIBLE_DEVICES=0,1,2,3 LOADWORKER=4 LIGHTLLM_TRITON_AUTOTUNE_LEVEL=1 \ + python -m lightllm.server.api_server \ + --model_dir /path/to/SenseNova-Loop \ + --model_name sensenova_loop \ + --mtp_mode eagle_with_att \ + --mtp_step 3 \ + --tp 4 \ + --port 8000 + +2P1D Deployment +--------------- + +This script uses eight GPUs to start two TP2 prefill services and one TP4 decode service, with MTP enabled: + +.. code-block:: bash + + MODEL_DIR=/path/to/SenseNova-Loop PORT=8000 \ + bash test/start_scripts/sensenova_loop/sensenova_loop_pd_2p1d.sh + +Set ``GPU_IDS`` to select the GPUs, or add ``DRY_RUN=1`` to preview the commands. +Press ``Ctrl+C`` to stop all services. + +Example Request +--------------- + +.. code-block:: bash + + curl http://localhost:8000/v1/chat/completions \ + -H 'Content-Type: application/json' \ + -d '{ + "model": "sensenova_loop", + "messages": [{"role": "user", "content": "Hello, introduce yourself in one sentence."}], + "max_tokens": 128 + }' diff --git a/docs/EN/source/index.rst b/docs/EN/source/index.rst index 0f5a142797..2b59db60b1 100755 --- a/docs/EN/source/index.rst +++ b/docs/EN/source/index.rst @@ -66,6 +66,7 @@ Documentation List GLM-4.7-Flash Deployment Qwen3.5 Deployment + SenseNova Loop Deployment .. toctree:: :maxdepth: 1 diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/stage1_single_token.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/stage1_single_token.py index d90aeb5bff..143b3c1ce1 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/stage1_single_token.py +++ b/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/stage1_single_token.py @@ -308,7 +308,7 @@ def mtp_diverse_stage1_single_token( Lq, Lk = int(q.shape[2]), int(k.shape[2]) assert Lq == Lk - assert Lk in {16, 32, 64, 128} + assert Lk in {16, 32, 64, 128, 256} batch = int(B_req_idx.shape[0]) kv_head_num = int(k.shape[1]) q_head_num = int(q.shape[1]) diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/stage2_single_token.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/stage2_single_token.py index b89b34ad14..5a56426932 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/stage2_single_token.py +++ b/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/stage2_single_token.py @@ -164,7 +164,7 @@ def mtp_diverse_stage2_single_token( num_stages = run_config["num_stages"] Lk = mid_out.shape[-1] - assert Lk in {16, 32, 64, 128} + assert Lk in {16, 32, 64, 128, 256} batch, head_num = mid_out.shape[0], mid_out.shape[1] grid = (batch, head_num) diff --git a/lightllm/common/req_manager/linear_att.py b/lightllm/common/req_manager/linear_att.py index c681203e3c..5ae9f4966c 100644 --- a/lightllm/common/req_manager/linear_att.py +++ b/lightllm/common/req_manager/linear_att.py @@ -90,10 +90,7 @@ def save_state(self, req: "InferReq", buffer_idx: int, state_cache_manager: Line dst_ssm_state.copy_(gpu_ssm_state, non_blocking=True) def get_mamba_cache(self, layer_idx_in_all: int): - assert ( - 0 <= layer_idx_in_all < self.linear_config.all_layer_num - ), f"invalid transformer layer index {layer_idx_in_all}" - layer_idx_in_linear = layer_idx_in_all - (layer_idx_in_all // self.linear_config.full_attention_interval) + layer_idx_in_linear = self.linear_config.get_linear_state_layer_index(layer_idx_in_all) conv_states = self.req_to_conv_state.buffer[layer_idx_in_linear] ssm_states = self.req_to_ssm_state.buffer[layer_idx_in_linear] return conv_states, ssm_states diff --git a/lightllm/common/state_cache_manager/__init__.py b/lightllm/common/state_cache_manager/__init__.py index 8e400e2d2c..758c6a88d4 100644 --- a/lightllm/common/state_cache_manager/__init__.py +++ b/lightllm/common/state_cache_manager/__init__.py @@ -17,6 +17,6 @@ def get_hybrid_cache_config(): return DeepseekV4CacheConfig.from_model_config(model_cfg, args) if model_type in ("glm5_next", "glm5_next_text"): return Glm5NextCacheConfig.from_model_config(model_cfg, args) - if model_type in ("qwen3_5", "qwen3_5_moe", "qwen3_5_text", "qwen3_5_moe_text"): + if model_type in ("qwen3_5", "sensenova_loop", "qwen3_5_moe", "qwen3_5_text", "qwen3_5_moe_text"): return LinearAttCacheConfig.load_from_args() raise ValueError("No hybrid state-cache layout registered for this model") diff --git a/lightllm/common/state_cache_manager/glm5_next.py b/lightllm/common/state_cache_manager/glm5_next.py index 13195a74ac..516552d7bc 100644 --- a/lightllm/common/state_cache_manager/glm5_next.py +++ b/lightllm/common/state_cache_manager/glm5_next.py @@ -21,10 +21,6 @@ def from_model_config(cls, config, args): linear = config["linear_attn_config"] tp = args.tp // args.dp assert linear["num_heads"] % tp == 0 - layers = config["num_hidden_layers"] - assert config["layer_types"] == [ - "deepseek_sparse_attention" if i % 4 == 3 else "linear_attention" for i in range(layers) - ] dtype = get_torch_dtype(args.data_type) # Raw index keys and compression scores only live in request tails. packed_dim = config["kv_lora_rank"] + cls.INDEX_PADDING_BYTES // dtype.itemsize @@ -41,11 +37,9 @@ def from_model_config(cls, config, args): head_linear_k_dim=linear["head_dim"], head_linear_v_dim=linear["head_dim"], conv_kernel_size=linear["short_conv_kernel_size"], - linear_layer_num=len(linear["kda_layers"]), conv_state_dtype=dtype, ssm_state_dtype=get_torch_dtype(args.linear_att_ssm_data_type), - full_attention_interval=4, - all_layer_num=layers, + layer_types=config["layer_types"], draft_full_att_kv_layer_num=get_added_mtp_kv_layer_num() if args.mtp_mode is not None else 0, index_kpool=config["index_kpool"], index_head_dim=config["index_head_dim"], diff --git a/lightllm/common/state_cache_manager/linear_att.py b/lightllm/common/state_cache_manager/linear_att.py index 35a37dcd28..496791d489 100644 --- a/lightllm/common/state_cache_manager/linear_att.py +++ b/lightllm/common/state_cache_manager/linear_att.py @@ -30,12 +30,32 @@ class LinearAttCacheConfig: head_linear_k_dim: int head_linear_v_dim: int conv_kernel_size: int - linear_layer_num: int conv_state_dtype: torch.dtype ssm_state_dtype: torch.dtype - full_attention_interval: int - all_layer_num: int # 包括 linear att 和 full att 的层加起来的层数 + layer_types: tuple[str, ...] draft_full_att_kv_layer_num: int = 0 + linear_layer_num: int = dataclasses.field(init=False) + all_layer_num: int = dataclasses.field(init=False) + + def __post_init__(self): + self.layer_types = tuple(self.layer_types) + indices = [] + full_count = linear_count = 0 + for kind in self.layer_types: + if kind == "linear_attention": + indices.append(linear_count) + linear_count += 1 + else: + indices.append(full_count) + full_count += 1 + # MTP draft layers only use full-attention KV slots after the main model. + indices.extend(range(full_count, full_count + self.draft_full_att_kv_layer_num)) + self._layer_indices = tuple(indices) + self.linear_layer_num = linear_count + self.all_layer_num = len(self.layer_types) + + def get_linear_state_layer_index(self, layer_index: int) -> int: + return self._layer_indices[layer_index] def get_conv_dim(self): # 第一项对应q的参数,第二项对应k的参数,第三项对应v的参数 @@ -49,7 +69,6 @@ def get_conv_dim(self): def get_main_model_full_att_layer_num(self): full_att_layer_num = self.all_layer_num - self.linear_layer_num - assert full_att_layer_num == self.all_layer_num // self.full_attention_interval return full_att_layer_num def get_full_att_kv_layer_num_with_draft_model(self): @@ -58,16 +77,7 @@ def get_full_att_kv_layer_num_with_draft_model(self): def get_full_att_kv_layer_index(self, layer_index: int) -> int: """Map a global target/draft layer index to the packed full-attention cache.""" - layer_index = int(layer_index) - if layer_index >= self.all_layer_num: - kv_layer_index = layer_index - self.linear_layer_num - else: - kv_layer_index = layer_index // self.full_attention_interval - assert 0 <= kv_layer_index < self.get_full_att_kv_layer_num_with_draft_model(), ( - f"layer {layer_index} maps outside the packed full-attention KV cache: " - f"slot={kv_layer_index}, slots={self.get_full_att_kv_layer_num_with_draft_model()}" - ) - return kv_layer_index + return self._layer_indices[layer_index] def get_conv_state_shape(self): # Base committed sliding-window state, without speculative MTP tail. @@ -117,17 +127,13 @@ def load_from_args() -> "LinearAttCacheConfig": from transformers.configuration_utils import PretrainedConfig model_cfg, _ = PretrainedConfig.get_config_dict(model_path) - model_type = model_cfg["model_type"] - assert model_type in ["qwen3_5", "qwen3_5_moe", "qwen3_5_text", "qwen3_5_moe_text"] - llm_config = model_cfg - try: - llm_config = llm_config["text_config"] - except: - pass - - n_layer = llm_config["num_hidden_layers"] + llm_config = model_cfg.get("text_config", model_cfg) + layer_types = llm_config["layer_types"] + if model_cfg["model_type"] == "sensenova_loop": + from lightllm.models.sensenova_loop.config import get_loop_layer_types - tp_world_size = get_env_start_args().tp // get_env_start_args().dp + layer_types = get_loop_layer_types(model_cfg["loop_settings"], llm_config) + tp_world_size = args.tp // args.dp full_att_dtype = ( torch.uint8 if args.llm_kv_type in {"fp8kv_sph", "fp8kv_spt"} else get_torch_dtype(args.data_type) ) @@ -144,11 +150,9 @@ def load_from_args() -> "LinearAttCacheConfig": head_linear_k_dim=llm_config["linear_key_head_dim"], head_linear_v_dim=llm_config["linear_value_head_dim"], conv_kernel_size=llm_config["linear_conv_kernel_dim"], - linear_layer_num=n_layer - (n_layer // llm_config["full_attention_interval"]), conv_state_dtype=get_torch_dtype(args.data_type), ssm_state_dtype=get_torch_dtype(args.linear_att_ssm_data_type), - full_attention_interval=llm_config["full_attention_interval"], - all_layer_num=n_layer, + layer_types=layer_types, draft_full_att_kv_layer_num=get_added_mtp_kv_layer_num(), ) diff --git a/lightllm/models/builtin.py b/lightllm/models/builtin.py index 0e0ca63505..d8e091d9aa 100644 --- a/lightllm/models/builtin.py +++ b/lightllm/models/builtin.py @@ -114,6 +114,9 @@ def _tarsier_text_model_is(model_type): ) ModelRegistry.register("qwen3", "lightllm.models.qwen3.model:Qwen3TpPartModel") ModelRegistry.register(["qwen3_5"], "lightllm.models.qwen3_5.model:Qwen3_5TpPartModel", is_multimodal=True) +ModelRegistry.register( + "sensenova_loop", "lightllm.models.sensenova_loop.model:SenseNovaLoopTpPartModel", is_multimodal=True +) ModelRegistry.register("qwen3_5_moe", "lightllm.models.qwen3_5_moe.model:Qwen3_5MOETpPartModel", is_multimodal=True) ModelRegistry.register("qwen3_moe", "lightllm.models.qwen3_moe.model:Qwen3MOEModel") ModelRegistry.register( @@ -186,7 +189,7 @@ def _tarsier_text_model_is(model_type): "lightllm.models.qwen3_5_moe_mtp.model:Qwen3_5MoeMTPModel", ) DraftModelRegistry.register( - ("qwen3_5", "qwen3_5_text"), + ("qwen3_5", "qwen3_5_text", "sensenova_loop"), ("vanilla_with_att", "eagle_with_att"), "lightllm.models.qwen3_5_mtp.model:Qwen3_5MTPModel", ) diff --git a/lightllm/models/qwen3next/model.py b/lightllm/models/qwen3next/model.py index 08233444eb..444ed1f393 100644 --- a/lightllm/models/qwen3next/model.py +++ b/lightllm/models/qwen3next/model.py @@ -11,7 +11,7 @@ ) from lightllm.models.qwen3next.infer_struct import Qwen3NextInferStateInfo from lightllm.utils.log_utils import init_logger -from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num, get_env_start_args +from lightllm.utils.envs_utils import get_env_start_args from lightllm.common.kv_cache_mem_manager.qwen3next_mem_manager import ( FP8StaticPerHeadQuantQwen3NextMemManager, FP8StaticPerTensorQuantQwen3NextMemManager, @@ -25,7 +25,6 @@ class Qwen3NextTpPartModel(Qwen3MOEModel): - # weight class pre_and_post_weight_class = Qwen3NextPreAndPostLayerWeight transformer_weight_class = Qwen3NextTransformerLayerWeight @@ -60,31 +59,6 @@ def _init_config(self): def _init_mem_manager(self): assert self.config["num_attention_heads"] % self.tp_world_size_ == 0 start_args: StartArgs = get_env_start_args() - ssm_dtype_dict = {"bfloat16": torch.bfloat16, "float32": torch.float32} - draft_full_att_kv_layer_num = get_added_mtp_kv_layer_num() - full_att_dtype = torch.uint8 if start_args.llm_kv_type in {"fp8kv_sph", "fp8kv_spt"} else self.data_type - self.linear_config = LinearAttCacheConfig( - tp_world_size=self.tp_world_size_, - full_att_all_num_kv_heads=self.config["num_key_value_heads"], - full_att_dtype=full_att_dtype, - full_att_num_kv_heads=self.num_kv_heads, - full_att_head_dim=self.config["head_dim"], - global_linear_k_heads=self.config["linear_num_key_heads"], - global_linear_v_heads=self.config["linear_num_value_heads"], - num_linear_k_heads=max(1, self.config["linear_num_key_heads"] // self.tp_world_size_), - num_linear_v_heads=max(1, self.config["linear_num_value_heads"] // self.tp_world_size_), - head_linear_k_dim=self.config["linear_key_head_dim"], - head_linear_v_dim=self.config["linear_value_head_dim"], - conv_kernel_size=self.config["linear_conv_kernel_dim"], - linear_layer_num=self.config["n_layer"] - - (self.config["n_layer"] // self.config["full_attention_interval"]), - conv_state_dtype=self.data_type, - ssm_state_dtype=ssm_dtype_dict[start_args.linear_att_ssm_data_type], - full_attention_interval=self.config["full_attention_interval"], - all_layer_num=self.config["n_layer"], - draft_full_att_kv_layer_num=draft_full_att_kv_layer_num, - ) - if start_args.llm_kv_type == "None": mem_manager_class = Qwen3NextMemManager elif start_args.llm_kv_type == "fp8kv_sph": @@ -104,6 +78,7 @@ def _init_mem_manager(self): ) def _init_req_manager(self): + self.linear_config = LinearAttCacheConfig.load_from_args() create_max_seq_len = 0 if self.batch_max_tokens is not None: @@ -112,7 +87,7 @@ def _init_req_manager(self): create_max_seq_len = max(create_max_seq_len, self.max_seq_length) self.req_manager = ReqManagerForMamba( - self.max_req_num, create_max_seq_len, None, linear_config=LinearAttCacheConfig.load_from_args() + self.max_req_num, create_max_seq_len, None, linear_config=self.linear_config ) return diff --git a/lightllm/models/sensenova_loop/__init__.py b/lightllm/models/sensenova_loop/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/models/sensenova_loop/config.py b/lightllm/models/sensenova_loop/config.py new file mode 100644 index 0000000000..0f5593b17d --- /dev/null +++ b/lightllm/models/sensenova_loop/config.py @@ -0,0 +1,14 @@ +"""SenseNova Loop execution layout.""" + + +def get_loop_layer_indices(settings: dict, num_layers: int) -> tuple[int, ...]: + """Map execution depth to checkpoint layers; caches use execution depth.""" + start, end, count = settings["start"], settings["end"], settings["eval_count"] + return tuple(range(start)) + tuple(range(start, end)) * count + tuple(range(end, num_layers)) + + +def get_loop_layer_types(settings: dict, text_config: dict) -> tuple[str, ...]: + num_layers = text_config["num_hidden_layers"] + indices = get_loop_layer_indices(settings, num_layers) + layer_types = text_config["layer_types"] + return tuple(layer_types[i] for i in indices) diff --git a/lightllm/models/sensenova_loop/layer_infer/__init__.py b/lightllm/models/sensenova_loop/layer_infer/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/models/sensenova_loop/layer_infer/transformer_layer_infer.py b/lightllm/models/sensenova_loop/layer_infer/transformer_layer_infer.py new file mode 100644 index 0000000000..688be0d866 --- /dev/null +++ b/lightllm/models/sensenova_loop/layer_infer/transformer_layer_infer.py @@ -0,0 +1,63 @@ +from lightllm.models.sensenova_loop.triton_kernel.mhc import hc_contract, hc_expand, hc_post, hc_pre_norm +from lightllm.models.sensenova_loop.triton_kernel.mhc.pre_norm import has_deepgemm_mhc +from lightllm.models.qwen3_5.layer_infer.transformer_layer_infer import Qwen35TransformerLayerInfer + + +class SenseNovaLoopTransformerLayerInfer(Qwen35TransformerLayerInfer): + def __init__(self, layer_num, network_config): + super().__init__(layer_num, network_config) + self.is_linear_attention_layer = network_config["layer_types"][layer_num] == "linear_attention" + self.physical_layer_num = network_config["loop_physical_indices"][layer_num] + self.mhc = network_config.get("mhc") + self.use_mhc = self.mhc is not None and self.mhc["start"] <= self.physical_layer_num < self.mhc["end"] + if not self.use_mhc: + return + self.use_deepgemm_mhc = self.mhc["streams"] == 4 and has_deepgemm_mhc() + self.mhc_entry = self.physical_layer_num == self.mhc["start"] + self.mhc_exit = self.physical_layer_num == self.mhc["end"] - 1 + + def context_forward(self, hidden, infer_state, layer_weight): + if not self.use_mhc: + return super().context_forward(hidden, infer_state, layer_weight) + return self._mhc_forward(hidden, infer_state, layer_weight, prefill=True) + + def token_forward(self, hidden, infer_state, layer_weight): + if not self.use_mhc: + return super().token_forward(hidden, infer_state, layer_weight) + return self._mhc_forward(hidden, infer_state, layer_weight, prefill=False) + + def _mhc_pre(self, residual, router_weight, norm_weight): + return hc_pre_norm( + x=residual, + fn=router_weight.weight, + base=router_weight.bias, + norm_weight=norm_weight.weight, + streams=self.mhc["streams"], + rms_eps=1e-6, + norm_eps=self.eps_, + sinkhorn_iters=self.mhc["iterations"], + use_deepgemm=self.use_deepgemm_mhc, + ) + + def _mhc_post(self, residual, branch, mix, post): + # SenseNova routers map output -> input; the common kernel uses input -> output. + return hc_post(branch, residual, mix.transpose(1, 2), post, self.mhc["streams"]) + + def _mhc_forward(self, hidden, infer_state, layer_weight, *, prefill): + streams = self.mhc["streams"] + if self.mhc_entry: + hidden = hc_expand(hidden, streams) + + branch, mix, post = self._mhc_pre(hidden, layer_weight.mhc_attn, layer_weight.att_norm_weight_) + if prefill: + branch = self.context_attention_forward(branch, infer_state, layer_weight) + else: + branch = self.token_attention_forward(branch, infer_state, layer_weight) + branch = branch.view(-1, self.embed_dim_) + hidden = self._mhc_post(hidden, branch, mix, post) + branch, mix, post = self._mhc_pre(hidden, layer_weight.mhc_ffn, layer_weight.ffn_norm_weight_) + branch = self._ffn(branch, infer_state, layer_weight).view(-1, self.embed_dim_) + hidden = self._mhc_post(hidden, branch, mix, post) + if self.mhc_exit: + return hc_contract(hidden, streams) + return hidden diff --git a/lightllm/models/sensenova_loop/layer_weights/__init__.py b/lightllm/models/sensenova_loop/layer_weights/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/models/sensenova_loop/layer_weights/transformer_layer_weight.py b/lightllm/models/sensenova_loop/layer_weights/transformer_layer_weight.py new file mode 100644 index 0000000000..f797ae914d --- /dev/null +++ b/lightllm/models/sensenova_loop/layer_weights/transformer_layer_weight.py @@ -0,0 +1,34 @@ +import torch + +from lightllm.common.basemodel.layer_weights.meta_weights import ParameterWeight +from lightllm.models.qwen3_5.layer_weights.transformer_layer_weight import Qwen35TransformerLayerWeight + + +class SenseNovaLoopTransformerLayerWeight(Qwen35TransformerLayerWeight): + def _init_weight(self): + execution_index = self.network_config_["loop_physical_indices"].index(self.layer_num_) + self.is_linear_attention_layer = self.network_config_["layer_types"][execution_index] == "linear_attention" + super()._init_weight() + self.mhc = self.network_config_.get("mhc") + if self.mhc is not None and self.mhc["start"] <= self.layer_num_ < self.mhc["end"]: + self._init_mhc() + + def _init_mhc(self): + prefix = f"model.layers.{self.layer_num_}" + streams = self.mhc["streams"] + mix_size = (2 + streams) * streams + flattened_hidden = streams * self.n_embed + self.mhc_attn = ParameterWeight( + weight_name=f"{prefix}.mhc_attn.weight", + bias_name=f"{prefix}.mhc_attn.bias", + data_type=torch.float32, + weight_shape=(mix_size, flattened_hidden), + bias_shape=(mix_size,), + ) + self.mhc_ffn = ParameterWeight( + weight_name=f"{prefix}.mhc_ffn.weight", + bias_name=f"{prefix}.mhc_ffn.bias", + data_type=torch.float32, + weight_shape=(mix_size, flattened_hidden), + bias_shape=(mix_size,), + ) diff --git a/lightllm/models/sensenova_loop/model.py b/lightllm/models/sensenova_loop/model.py new file mode 100644 index 0000000000..ce6d36a7d5 --- /dev/null +++ b/lightllm/models/sensenova_loop/model.py @@ -0,0 +1,51 @@ +"""SenseNova Loop: shared checkpoint weights and independent execution caches.""" + +from lightllm.models.qwen3_5.model import Qwen3_5TpPartModel +from lightllm.utils.config_utils import get_config_json +from .config import get_loop_layer_indices, get_loop_layer_types +from .layer_infer.transformer_layer_infer import SenseNovaLoopTransformerLayerInfer +from .layer_weights.transformer_layer_weight import SenseNovaLoopTransformerLayerWeight + + +class SenseNovaLoopTpPartModel(Qwen3_5TpPartModel): + """Share physical weights, with separate KV/conv/SSM slots per execution layer.""" + + transformer_weight_class = SenseNovaLoopTransformerLayerWeight + transformer_layer_infer_class = SenseNovaLoopTransformerLayerInfer + + def _init_config(self): + super()._init_config() + model_config = get_config_json(self.weight_dir_) + settings = model_config["loop_settings"] + self.num_physical_layers = self.config["n_layer"] + self.loop_layer_indices = get_loop_layer_indices(settings, self.num_physical_layers) + self.config.update( + layer_types=get_loop_layer_types(settings, self.config), + n_layer=len(self.loop_layer_indices), + num_hidden_layers=len(self.loop_layer_indices), + loop_physical_indices=self.loop_layer_indices, + ) + if model_config.get("mhc") is not None: + self.config["mhc"] = model_config["mhc"] + + def _init_weights(self, start_layer_index=0): + # Allocate and preprocess each checkpoint layer once, then alias weights. + self.pre_post_weight = self.pre_and_post_weight_class(self.data_type, network_config=self.config) + self.physical_layer_weights = [ + self.transformer_weight_class( + i, + self.data_type, + network_config=self.config, + quant_cfg=self.quant_cfg, + ) + for i in range(start_layer_index, start_layer_index + self.num_physical_layers) + ] + self.trans_layers_weight = [self.physical_layer_weights[i] for i in self.loop_layer_indices] + + def load_weights(self, weight_dict): + execution_weights = self.trans_layers_weight + self.trans_layers_weight = self.physical_layer_weights + try: + super().load_weights(weight_dict) + finally: + self.trans_layers_weight = execution_weights diff --git a/lightllm/models/sensenova_loop/triton_kernel/__init__.py b/lightllm/models/sensenova_loop/triton_kernel/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/models/sensenova_loop/triton_kernel/mhc/__init__.py b/lightllm/models/sensenova_loop/triton_kernel/mhc/__init__.py new file mode 100644 index 0000000000..79e1ba0ecc --- /dev/null +++ b/lightllm/models/sensenova_loop/triton_kernel/mhc/__init__.py @@ -0,0 +1,7 @@ +"""SenseNova Loop mHC routing, preserving checkpoint rounding semantics.""" + +from .post import hc_post +from .pre_norm import hc_pre_norm +from .streams import hc_contract, hc_expand + +__all__ = ["hc_contract", "hc_expand", "hc_post", "hc_pre_norm"] diff --git a/lightllm/models/sensenova_loop/triton_kernel/mhc/post.py b/lightllm/models/sensenova_loop/triton_kernel/mhc/post.py new file mode 100644 index 0000000000..7fc7905e7a --- /dev/null +++ b/lightllm/models/sensenova_loop/triton_kernel/mhc/post.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""mHC post-mixing of sublayer outputs into residual streams.""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _hc_post_kernel( + layer_output, # 输入 [T, H],激活 dtype。 + residual, # 输入 [T, S * H],激活 dtype;当前子层 pre-mixing 前的残差。 + residual_mix, # 输入 [T, S, S],fp32;最后两维为输入、输出 stream。 + post_mix, # 输入 [T, S],fp32。 + output, # 输出 [T, S * H],与 layer_output 相同 dtype。 + hidden: tl.constexpr, + layer_stride_m: tl.constexpr, + residual_stride_m: tl.constexpr, + mix_stride_m: tl.constexpr, + mix_stride_i: tl.constexpr, + mix_stride_j: tl.constexpr, + post_stride_m: tl.constexpr, + out_stride_m: tl.constexpr, + BLOCK_H: tl.constexpr, + STREAMS: tl.constexpr, +): + """将子层输出和原始残差混合为两条或四条新 stream。 + + T = tokens,S = STREAMS,H = hidden。每个 program 同时计算 S 条 stream 的 BLOCK_H 个 + hidden 元素,沿 residual_mix 的输入 stream 维归约;fp32 累加后写入 output。 + 输出按 [T, S * H] 展平存储,不返回 Tensor。 + """ + token = tl.program_id(0) + hidden_block = tl.program_id(1) + hidden_offsets = hidden_block * BLOCK_H + tl.arange(0, BLOCK_H) + hidden_mask = hidden_offsets < hidden + + # Compute all stream outputs in one program so the layer output and + # residual streams are read only once, reducing prefill memory traffic. + layer_value = tl.load( + layer_output + token * layer_stride_m + hidden_offsets, + mask=hidden_mask, + other=0.0, + ).to(tl.float32) + residual_base = residual + token * residual_stride_m + hidden_offsets + residual_0 = tl.load(residual_base, mask=hidden_mask, other=0.0).to(tl.float32) + residual_1 = tl.load(residual_base + hidden, mask=hidden_mask, other=0.0).to(tl.float32) + if STREAMS == 4: + residual_2 = tl.load(residual_base + 2 * hidden, mask=hidden_mask, other=0.0).to(tl.float32) + residual_3 = tl.load(residual_base + 3 * hidden, mask=hidden_mask, other=0.0).to(tl.float32) + + post_base = post_mix + token * post_stride_m + mix_base = residual_mix + token * mix_stride_m + for stream in tl.static_range(STREAMS): + column = mix_base + stream * mix_stride_j + mix_0 = tl.load(column) + mix_1 = tl.load(column + mix_stride_i) + if STREAMS == 4: + mix_2 = tl.load(column + 2 * mix_stride_i) + mix_3 = tl.load(column + 3 * mix_stride_i) + post_value = tl.load(post_base + stream) + mix_0 = mix_0.to(layer_output.dtype.element_ty).to(tl.float32) + mix_1 = mix_1.to(layer_output.dtype.element_ty).to(tl.float32) + if STREAMS == 4: + mix_2 = mix_2.to(layer_output.dtype.element_ty).to(tl.float32) + mix_3 = mix_3.to(layer_output.dtype.element_ty).to(tl.float32) + post_value = post_value.to(layer_output.dtype.element_ty).to(tl.float32) + contribution = layer_value * post_value + accumulator = tl.zeros((BLOCK_H,), tl.float32) + accumulator += residual_0 * mix_0 + accumulator += residual_1 * mix_1 + if STREAMS == 4: + accumulator += residual_2 * mix_2 + accumulator += residual_3 * mix_3 + accumulator = accumulator.to(layer_output.dtype.element_ty).to(tl.float32) + contribution = contribution.to(layer_output.dtype.element_ty).to(tl.float32) + accumulator += contribution + tl.store(output + token * out_stride_m + stream * hidden + hidden_offsets, accumulator, mask=hidden_mask) + + +def hc_post( + layer_output: torch.Tensor, + residual: torch.Tensor, + residual_mix: torch.Tensor, + post_mix: torch.Tensor, + streams: int, +) -> torch.Tensor: + """用一次 Triton launch 将子层输出混回残差 streams。 + + T = tokens,S = streams,H = 单条 stream 的 hidden size。 + + Args: + layer_output: [T, H],连续的 attention / FFN 子层输出。 + residual: [T, S * H],连续的、该子层 pre-mixing 前保存的原始残差。 + residual_mix: [T, S, S],fp32;来自 hc_pre_norm。 + post_mix: [T, S],fp32;来自同一次 pre 调用。 + streams: 残差 stream 数 S;支持 S = 2 或 4。 + + Returns: + 新残差 [T, S * H],dtype 与 layer_output 相同;输入张量保持不变。 + 逻辑输出 [T, S, H] 以 fp32 计算: + out[t, j, h] = post_mix[t, j] * layer_output[t, h] + + sum_i residual_mix[t, i, j] * residual.view(T, S, H)[t, i, h]。 + 矩阵第一个 stream 维 i 是输入,第二个 stream 维 j 是输出。 + """ + + tokens, hidden = layer_output.shape + assert streams in (2, 4) + assert layer_output.is_contiguous() and residual.is_contiguous() + output = torch.empty( + (tokens, streams * hidden), + dtype=layer_output.dtype, + device=layer_output.device, + ) + block_h = min(triton.next_power_of_2(hidden), 1024) + _hc_post_kernel[(tokens, triton.cdiv(hidden, block_h))]( + layer_output, + residual, + residual_mix, + post_mix, + output, + hidden=hidden, + layer_stride_m=layer_output.stride(0), + residual_stride_m=residual.stride(0), + mix_stride_m=residual_mix.stride(0), + mix_stride_i=residual_mix.stride(1), + mix_stride_j=residual_mix.stride(2), + post_stride_m=post_mix.stride(0), + out_stride_m=output.stride(0), + BLOCK_H=block_h, + STREAMS=streams, + num_warps=8, + ) + return output diff --git a/lightllm/models/sensenova_loop/triton_kernel/mhc/pre_norm.py b/lightllm/models/sensenova_loop/triton_kernel/mhc/pre_norm.py new file mode 100644 index 0000000000..522dd0d15b --- /dev/null +++ b/lightllm/models/sensenova_loop/triton_kernel/mhc/pre_norm.py @@ -0,0 +1,251 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""mHC pre-mixing fused with the following RMSNorm.""" + +from __future__ import annotations + +import os +from functools import lru_cache +from typing import Tuple + +import torch +import triton +import triton.language as tl + +LIGHTLLM_DISABLE_DEEPGEMM_MHC = os.getenv("LIGHTLLM_DISABLE_DEEPGEMM_MHC", "False").upper() in ["ON", "TRUE", "1"] + + +@lru_cache(maxsize=1) +def has_deepgemm_mhc() -> bool: + try: + from deep_gemm import tf32_hc_prenorm_gemm + except ImportError: + return False + return True + + +@triton.jit +def _hc_prepare_prenorm_kernel( + gemm_partial, # 输入 [P, T, M],fp32;投影的 split-K 部分和。 + sqrsum_partial, # 输入 [P, T],fp32;残差平方和的 split-K 部分和。 + base, # 输入 [M],fp32;M = 2 * S + S * S。 + pre, # 输出 [T, S],fp32。 + post, # 输出 [T, S],fp32。 + residual_mix, # 输出 [T, S, S],fp32;最后两维为输入、输出 stream。 + gemm_stride_s, + gemm_stride_m: tl.constexpr, + sqrsum_stride_s, + sqrsum_stride_m: tl.constexpr, + pre_stride_m: tl.constexpr, + residual_stride_m: tl.constexpr, + FLATTENED_HIDDEN: tl.constexpr, + RMS_EPS: tl.constexpr, + STREAMS: tl.constexpr, + SINKHORN_ITERS: tl.constexpr, + N_SPLITS: tl.constexpr, +): + """归并分片并准备混合权重:[P, T, M] / [P, T] -> [T, S] / [T, S, S]。 + + T = tokens,S = STREAMS,M = 2 * S + S * S,P = N_SPLITS。 + FLATTENED_HIDDEN = S * H,用于从平方和计算每个 token 的 RMS 倒数。 + 每个 program 处理一个 token,写入 pre/post/residual_mix,不返回 Tensor。 + """ + token = tl.program_id(0) + stream_offsets = tl.arange(0, STREAMS) + matrix_offsets = tl.arange(0, STREAMS * STREAMS) + pre_raw = tl.zeros((STREAMS,), dtype=tl.float32) + post_raw = tl.zeros((STREAMS,), dtype=tl.float32) + matrix_raw = tl.zeros((STREAMS * STREAMS,), dtype=tl.float32) + sqrsum = 0.0 + for split in tl.static_range(N_SPLITS): + partial_base = gemm_partial + split * gemm_stride_s + token * gemm_stride_m + pre_raw += tl.load(partial_base + stream_offsets) + post_raw += tl.load(partial_base + STREAMS + stream_offsets) + matrix_raw += tl.load(partial_base + 2 * STREAMS + matrix_offsets) + sqrsum += tl.load(sqrsum_partial + split * sqrsum_stride_s + token * sqrsum_stride_m) + inv_rms = tl.rsqrt(sqrsum / FLATTENED_HIDDEN + RMS_EPS) + pre_raw *= inv_rms + post_raw *= inv_rms + matrix_raw *= inv_rms + + pre_values = tl.sigmoid(pre_raw + tl.load(base + stream_offsets)) + post_values = 2.0 * tl.sigmoid(post_raw + tl.load(base + STREAMS + stream_offsets)) + + residual_logits = matrix_raw + tl.load(base + 2 * STREAMS + matrix_offsets) + residual_logits = tl.reshape(residual_logits, (STREAMS, STREAMS)) + for _ in range(SINKHORN_ITERS): + residual_logits -= tl.max(residual_logits, axis=1)[:, None] + residual_logits -= tl.log(tl.sum(tl.exp(residual_logits), axis=1))[:, None] + residual_logits -= tl.max(residual_logits, axis=0)[None, :] + residual_logits -= tl.log(tl.sum(tl.exp(residual_logits), axis=0))[None, :] + matrix = tl.exp(residual_logits) + + tl.store(pre + token * pre_stride_m + stream_offsets, pre_values) + tl.store(post + token * pre_stride_m + stream_offsets, post_values) + tl.store( + residual_mix + token * residual_stride_m + matrix_offsets, + tl.reshape(matrix, (STREAMS * STREAMS,)), + ) + + +@triton.jit +def _hc_pre_combine_norm_kernel( + x, # 输入 [T, S * H],bf16;逻辑形状 [T, S, H]。 + pre, # 输入 [T, S],fp32。 + norm_weight, # 输入 [H],子层 RMSNorm 权重,通常为 bf16。 + output, # 输出 [T, H],bf16;加权合并后再做 RMSNorm。 + hidden: tl.constexpr, + x_stride_m: tl.constexpr, + pre_stride_m: tl.constexpr, + out_stride_m: tl.constexpr, + STREAMS: tl.constexpr, + NORM_EPS: tl.constexpr, + BLOCK_H: tl.constexpr, +): + """融合 stream 合并与子层 RMSNorm:[T, S, H] -> [T, H]。 + + T = tokens,S = STREAMS,H = hidden。每个 program 处理一个 token 的 + 整行 hidden;先把合并结果舍入到 bf16,再以 fp32 计算均方值、归一化并乘 norm_weight。 + 结果写入 output,不返回 Tensor。 + """ + token = tl.program_id(0) + hidden_offsets = tl.arange(0, BLOCK_H) + hidden_mask = hidden_offsets < hidden + accumulator = tl.zeros((BLOCK_H,), dtype=tl.float32) + for stream in tl.static_range(STREAMS): + residual = tl.load( + x + token * x_stride_m + stream * hidden + hidden_offsets, + mask=hidden_mask, + other=0.0, + ).to(tl.float32) + pre_value = tl.load(pre + token * pre_stride_m + stream) + pre_value = pre_value.to(x.dtype.element_ty).to(tl.float32) + accumulator += residual * pre_value + + # Preserve the checkpoint's bf16 rounding between stream mixing and RMSNorm. + rounded = accumulator.to(x.dtype.element_ty).to(tl.float32) + variance = tl.sum(rounded * rounded, axis=0) / hidden + inv_rms = tl.rsqrt(variance + NORM_EPS) + weight = tl.load(norm_weight + hidden_offsets, mask=hidden_mask, other=0.0) + tl.store( + output + token * out_stride_m + hidden_offsets, + rounded * inv_rms * weight, + mask=hidden_mask, + ) + + +def _compute_prenorm_splits(tokens: int, flattened_hidden: int, device: torch.device) -> int: + """由 T = tokens、K = flattened_hidden 和设备 SM 数确定 split-K 分片数 P。 + + tokens 和 flattened_hidden 为整数,device 指定 CUDA 设备;返回整数 P。 + P 决定 gemm_partial [P, T, M] 和 sqrsum_partial [P, T] 的首维, + 其中 M = 2 * S + S * S,S 为 stream 数。 + """ + grid_size = triton.cdiv(tokens, 64) + k_blocks = triton.cdiv(flattened_hidden, 64) + sms = torch.cuda.get_device_properties(device).multi_processor_count + return max(1, min(sms // max(grid_size, 1), k_blocks // 4)) + + +def hc_pre_norm( + x: torch.Tensor, + fn: torch.Tensor, + base: torch.Tensor, + norm_weight: torch.Tensor, + streams: int, + rms_eps: float, + norm_eps: float, + sinkhorn_iters: int, + *, + use_deepgemm: bool = True, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """将 mHC pre-mixing 与紧随其后的子层 RMSNorm 融合。 + + T = tokens,S = streams,H = 单条 stream 的 hidden size,M = 2 * S + S * S。 + + Args: + x: [T, S * H],连续的残差激活;逻辑形状 [T, S, H]。 + fn: [M, S * H],连续的 fp32 投影权重。 + base: [M],fp32;按 [S, S, S * S] 分组的偏置。 + norm_weight: [H],连续的子层 RMSNorm 权重,通常为 bf16。 + streams: 残差 stream 数 S;支持 S = 2 或 4。 + rms_eps: 在 S * H 维计算投影前 RMS 倒数时使用的 epsilon。 + norm_eps: 合并 stream 后,在 H 维做子层 RMSNorm 时使用的 epsilon。 + sinkhorn_iters: Sinkhorn 迭代次数。 + use_deepgemm: 四 stream 可使用 DeepGEMM 投影;两 stream 或 false 使用 PyTorch 投影。 + 两条路径共用融合的 Triton Sinkhorn / stream 合并 / RMSNorm。 + + Returns: + (layer_input, residual_mix, post_mix):分别为 bf16 [T, H]、fp32 [T, S, S]、fp32 [T, S]。 + layer_input 已完成子层 RMSNorm;residual_mix[t, i, j] 对应输入 i 到输出 j。 + x 保持不变,留给对应 hc_post 使用。 + + 默认用 DeepGEMM 融合投影与平方和计算,Triton 完成 Sinkhorn、stream 合并和 RMSNorm。 + 合并结果先舍入到 bf16,再做 RMSNorm。 + LIGHTLLM_DISABLE_DEEPGEMM_MHC=1 禁用 DeepGEMM,改用 PyTorch 计算投影与平方和。 + """ + + assert x.ndim == 2 and x.shape[-1] % streams == 0 + assert streams in (2, 4) + assert x.dtype in (torch.bfloat16, torch.float16, torch.float32) and fn.dtype == torch.float32 + assert x.is_contiguous() and fn.is_contiguous() and norm_weight.is_contiguous() + assert fn.shape == (2 * streams + streams * streams, x.shape[1]) + tokens, flattened_hidden = x.shape + hidden = flattened_hidden // streams + + if streams == 2 or not use_deepgemm or LIGHTLLM_DISABLE_DEEPGEMM_MHC or x.dtype != torch.bfloat16: + # PyTorch produces one complete partition: [1, T, M] and [1, T]. + x_fp32 = x.float() + gemm_partial = (x_fp32 @ fn.T).unsqueeze(0) + sqrsum_partial = x_fp32.square().sum(dim=-1).unsqueeze(0) + n_splits = 1 + else: + from deep_gemm import tf32_hc_prenorm_gemm + + mix_size = (2 + streams) * streams + n_splits = _compute_prenorm_splits(tokens, flattened_hidden, x.device) + gemm_partial = torch.empty((n_splits, tokens, mix_size), dtype=torch.float32, device=x.device) + sqrsum_partial = torch.empty((n_splits, tokens), dtype=torch.float32, device=x.device) + tf32_hc_prenorm_gemm(x, fn, gemm_partial, sqrsum_partial, n_splits) + + pre = torch.empty((tokens, streams), dtype=torch.float32, device=x.device) + post = torch.empty_like(pre) + residual_mix = torch.empty((tokens, streams, streams), dtype=torch.float32, device=x.device) + _hc_prepare_prenorm_kernel[(tokens,)]( + gemm_partial, + sqrsum_partial, + base, + pre, + post, + residual_mix, + gemm_partial.stride(0), + gemm_partial.stride(1), + sqrsum_partial.stride(0), + sqrsum_partial.stride(1), + pre.stride(0), + residual_mix.stride(0), + FLATTENED_HIDDEN=flattened_hidden, + RMS_EPS=rms_eps, + STREAMS=streams, + SINKHORN_ITERS=sinkhorn_iters, + N_SPLITS=n_splits, + num_warps=1, + ) + + layer_input = torch.empty((tokens, hidden), dtype=x.dtype, device=x.device) + block_h = triton.next_power_of_2(hidden) + _hc_pre_combine_norm_kernel[(tokens,)]( + x, + pre, + norm_weight, + layer_input, + hidden=hidden, + x_stride_m=x.stride(0), + pre_stride_m=pre.stride(0), + out_stride_m=layer_input.stride(0), + STREAMS=streams, + NORM_EPS=norm_eps, + BLOCK_H=block_h, + num_warps=8, + ) + return layer_input, residual_mix, post diff --git a/lightllm/models/sensenova_loop/triton_kernel/mhc/streams.py b/lightllm/models/sensenova_loop/triton_kernel/mhc/streams.py new file mode 100644 index 0000000000..4c2c56757c --- /dev/null +++ b/lightllm/models/sensenova_loop/triton_kernel/mhc/streams.py @@ -0,0 +1,31 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Expansion and mean contraction of flattened residual streams.""" + +from __future__ import annotations + +import torch + + +def hc_expand(x: torch.Tensor, streams: int) -> torch.Tensor: + """将每个 token 的 embedding 复制到 S 条 residual stream。 + + 输入 x: [T, H],streams: 标量 S;T = tokens,H = hidden size。 + 输出: [T, S * H],dtype/device 与 x 相同,逻辑形状为 [T, S, H]。 + output.view(T, S, H)[t, s, h] = x[t, h],沿 stream 维复制完整的 hidden 向量。 + """ + + assert x.ndim == 2 + return x.unsqueeze(1).expand(-1, streams, -1).reshape(x.shape[0], -1) + + +def hc_contract(x: torch.Tensor, streams: int) -> torch.Tensor: + """沿 stream 维取均值,将模型末尾的多条残差收缩成一个 hidden 向量。 + + 输入 x: [T, S * H],streams: 标量 S;T = tokens,H = 单条 stream 的 hidden size。 + 输出: [T, H],dtype/device 与 x 相同,即 x.view(T, S, H).mean(dim=1)。 + 此操作不含可学习权重;模型若使用 hc_head,应调用其专用的加权收缩算子。 + """ + + assert x.ndim == 2 and x.shape[-1] % streams == 0 + return x.view(x.shape[0], streams, -1).mean(dim=1) diff --git a/lightllm/server/tokenizer.py b/lightllm/server/tokenizer.py index 71c36ee55e..4b5ae89530 100644 --- a/lightllm/server/tokenizer.py +++ b/lightllm/server/tokenizer.py @@ -126,14 +126,12 @@ def get_tokenizer( tokenizer = QWen3VLTokenizer( tokenizer=tokenizer, image_processor=processor.image_processor, model_cfg=model_cfg ) - elif model_type in ["qwen3_5", "qwen3_5_moe"] and "vision_config" in model_cfg: - from transformers import AutoProcessor + elif model_type in ["qwen3_5", "sensenova_loop", "qwen3_5_moe"] and "vision_config" in model_cfg: from ..models.qwen3_5.model import QWen3_5Tokenizer + from transformers import AutoImageProcessor - processor = AutoProcessor.from_pretrained(tokenizer_name) - tokenizer = QWen3_5Tokenizer( - tokenizer=tokenizer, image_processor=processor.image_processor, model_cfg=model_cfg - ) + image_processor = AutoImageProcessor.from_pretrained(tokenizer_name) + tokenizer = QWen3_5Tokenizer(tokenizer=tokenizer, image_processor=image_processor, model_cfg=model_cfg) elif model_cfg.get("thinker_config") is not None: from ..models.qwen3_omni_moe_thinker.model import QWen3OmniTokenizer from transformers import AutoProcessor diff --git a/lightllm/server/visualserver/model_infer/model_rpc.py b/lightllm/server/visualserver/model_infer/model_rpc.py index 6c6e1079eb..6451e70008 100644 --- a/lightllm/server/visualserver/model_infer/model_rpc.py +++ b/lightllm/server/visualserver/model_infer/model_rpc.py @@ -92,7 +92,7 @@ def exposed_init_model(self, kvargs): self.model = ( Qwen2_5_VisionTransformerPretrainedModel(kvargs, **model_cfg["vision_config"]).eval().bfloat16() ) - elif self.model_type in ["qwen3_vl", "qwen3_vl_moe", "qwen3_5", "qwen3_5_moe"]: + elif self.model_type in ["qwen3_vl", "qwen3_vl_moe", "qwen3_5", "sensenova_loop", "qwen3_5_moe"]: self.model = ( Qwen3VisionTransformerPretrainedModel(kvargs, **model_cfg["vision_config"]).eval().bfloat16() ) diff --git a/lightllm/utils/config_utils.py b/lightllm/utils/config_utils.py index 2cdbc07dce..7b258959e9 100644 --- a/lightllm/utils/config_utils.py +++ b/lightllm/utils/config_utils.py @@ -323,7 +323,7 @@ def get_eos_token_ids(model_path: str) -> Optional[List[int]]: try: config_json = get_config_json(model_path) model_type = config_json.get("model_type") or config_json.get("text_config", {}).get("model_type") - if model_type in {"qwen3_5", "qwen3_5_text", "qwen3_5_moe", "qwen3_5_moe_text"}: + if model_type in {"qwen3_5", "sensenova_loop", "qwen3_5_text", "qwen3_5_moe", "qwen3_5_moe_text"}: from transformers import AutoTokenizer eos_token_ids = [] @@ -467,7 +467,7 @@ def has_vision_module(model_path: str) -> bool: return True elif model_type == "neo_chat": return True - elif model_type in ["qwen3_5", "qwen3_5_moe"]: + elif model_type in ["qwen3_5", "sensenova_loop", "qwen3_5_moe"]: return True else: raise Exception("unknown vision model type") @@ -506,10 +506,15 @@ def is_linear_att_mixed_model(model_path: str) -> bool: model_cfg, _ = PretrainedConfig.get_config_dict(model_path) model_type = model_cfg["model_type"] - if model_type in ["qwen3_5", "qwen3_5_moe", "qwen3_5_text", "qwen3_5_moe_text", "glm5_next", "glm5_next_text"]: - return True - else: - return False + return model_type in [ + "qwen3_5", + "sensenova_loop", + "qwen3_5_moe", + "qwen3_5_text", + "qwen3_5_moe_text", + "glm5_next", + "glm5_next_text", + ] except: logger.info(f"model path: {model_path} does not has linear hybrid attention") return False @@ -548,7 +553,7 @@ def get_tool_call_parser_for_model(model_path: str) -> Optional[str]: return None # Qwen3.5 series - if model_type in ["qwen3_5", "qwen3_5_moe", "qwen3_5_text", "qwen3_5_moe_text"]: + if model_type in ["qwen3_5", "sensenova_loop", "qwen3_5_moe", "qwen3_5_text", "qwen3_5_moe_text"]: return "qwen3_coder" # Qwen3 series @@ -598,6 +603,7 @@ def get_reasoning_parser_for_model(model_path: str) -> Optional[str]: "qwen3_vl_moe_text", "qwen3_omni_moe", "qwen3_5", + "sensenova_loop", "qwen3_5_moe", "qwen3_5_text", "qwen3_5_moe_text", diff --git a/test/cpu_cache_kernel/test_speed.py b/test/cpu_cache_kernel/test_speed.py index f187088883..d59142e09f 100644 --- a/test/cpu_cache_kernel/test_speed.py +++ b/test/cpu_cache_kernel/test_speed.py @@ -6,9 +6,9 @@ full_att_num_kv_heads=1, full_att_head_dim=256, num_linear_k_heads=2, num_linear_v_heads=8, head_linear_k_dim=128, head_linear_v_dim=128, - conv_kernel_size=4, linear_layer_num=36, + conv_kernel_size=4, 36 linear layers and 12 full-attention layers, conv_state_dtype=torch.bfloat16, ssm_state_dtype=torch.float32, - full_attention_interval=4, all_layer_num=48 + layer_types=(linear_attention, linear_attention, linear_attention, full_attention) * 12 """ import os @@ -47,16 +47,16 @@ full_att_dtype=torch.bfloat16, full_att_num_kv_heads=1, full_att_head_dim=256, + global_linear_k_heads=16, + global_linear_v_heads=64, num_linear_k_heads=2, num_linear_v_heads=8, head_linear_k_dim=128, head_linear_v_dim=128, conv_kernel_size=4, - linear_layer_num=36, conv_state_dtype=torch.bfloat16, ssm_state_dtype=torch.float32, - full_attention_interval=4, - all_layer_num=48, + layer_types=("linear_attention", "linear_attention", "linear_attention", "full_attention") * 12, ) print(f"LinearAttCacheConfig:\n{linear_config}\n", flush=True) @@ -64,7 +64,7 @@ # Step 2 – derive sizes from the config # --------------------------------------------------------------------------- big_page_token_num = _env_args["cpu_cache_token_page_size"] # 512 -full_att_layer_num = linear_config.all_layer_num // linear_config.full_attention_interval # 12 +full_att_layer_num = linear_config.get_main_model_full_att_layer_num() # 12 full_att_bytes = linear_config.get_cpu_cache_full_att_bytes() # per big page conv_bytes = linear_config.get_cpu_cache_conv_bytes() diff --git a/test/start_scripts/sensenova_loop/sensenova_loop_mtp.sh b/test/start_scripts/sensenova_loop/sensenova_loop_mtp.sh new file mode 100644 index 0000000000..96751f64ab --- /dev/null +++ b/test/start_scripts/sensenova_loop/sensenova_loop_mtp.sh @@ -0,0 +1,33 @@ +#!/usr/bin/env bash +set -euo pipefail + +if [[ "$#" -lt 2 ]]; then + echo "Usage: $0 [server arguments...]" >&2 + exit 2 +fi + +MODEL_DIR="$1" +MTP_DIR="$2" +shift 2 + +export LOADWORKER="${LOADWORKER:-4}" +export LIGHTLLM_TRITON_AUTOTUNE_LEVEL="${LIGHTLLM_TRITON_AUTOTUNE_LEVEL:-1}" + +# A single MTP layer drafts recurrently, so mtp_step can exceed its layer count. +exec python -m lightllm.server.api_server \ + --model_dir "${MODEL_DIR}" \ + --mtp_draft_model_dir "${MTP_DIR}" \ + --mtp_mode eagle_with_att \ + --mtp_step "${MTP_STEP:-3}" \ + --tp "${TP:-4}" \ + --data_type bfloat16 \ + --llm_prefill_att_backend fa3 \ + --llm_decode_att_backend fa3 \ + --max_req_total_len 32768 \ + --batch_max_tokens 8192 \ + --running_max_req_size 32 \ + --graph_max_batch_size 32 \ + --linear_att_cache_size 64 \ + --reasoning_parser qwen3 \ + --port "${PORT:-8000}" \ + "$@" diff --git a/test/start_scripts/sensenova_loop/sensenova_loop_pd_2p1d.sh b/test/start_scripts/sensenova_loop/sensenova_loop_pd_2p1d.sh new file mode 100644 index 0000000000..fd2f4a147a --- /dev/null +++ b/test/start_scripts/sensenova_loop/sensenova_loop_pd_2p1d.sh @@ -0,0 +1,94 @@ +#!/usr/bin/env bash +# Two Prefill services and one Decode service; TP defaults to 2 / 4. +# Set MODEL_DIR and GPU_IDS as needed. DRY_RUN=1 prints the startup commands. +set -euo pipefail + +LIGHTLLM_ROOT="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/../../.." && pwd)" +MODEL_DIR="${MODEL_DIR:-/mnt/afs_agents/yankuihao/model_ckpts/SenseNova-6.8-Loop-9B-Preview-mtp13200}" +PORT="${PORT:-8000}" +PREFILL_TP="${PREFILL_TP:-2}" +DECODE_TP="${DECODE_TP:-4}" +GPU_IDS="${GPU_IDS:-0,1,2,3,4,5,6,7}" +DRY_RUN="${DRY_RUN:-0}" + +CHAT_TEMPLATE_ARGS=() +if [[ -n "${CHAT_TEMPLATE:-}" ]]; then + CHAT_TEMPLATE_ARGS=(--chat_template "$(realpath -- "$CHAT_TEMPLATE")") +fi + +export LOADWORKER="${LOADWORKER:-8}" LIGHTLLM_TRITON_AUTOTUNE_LEVEL="${LIGHTLLM_TRITON_AUTOTUNE_LEVEL:-1}" +export PYTHONUNBUFFERED=1 +export NO_PROXY="${NO_PROXY:+${NO_PROXY},}127.0.0.1,localhost" +export no_proxy="$NO_PROXY" + +IFS=',' read -r -a GPUS <<< "$GPU_IDS" +if (( ${#GPUS[@]} != 2 * PREFILL_TP + DECODE_TP )); then + printf 'GPU_IDS must contain 2*PREFILL_TP+DECODE_TP GPU IDs.\n' >&2 + exit 2 +fi +join_gpus() { local IFS=','; printf '%s' "$*"; } +P1_GPUS="$(join_gpus "${GPUS[@]:0:PREFILL_TP}")" +P2_GPUS="$(join_gpus "${GPUS[@]:PREFILL_TP:PREFILL_TP}")" +D_GPUS="$(join_gpus "${GPUS[@]:$((2 * PREFILL_TP)):DECODE_TP}")" + +MODEL_ARGS=( + --model_dir "$MODEL_DIR" --model_name "${MODEL_NAME:-sensenova_loop}" + --host 0.0.0.0 --max_image_token_count 4096 "${CHAT_TEMPLATE_ARGS[@]}" +) +COMMON_ARGS=( + "${MODEL_ARGS[@]}" --mtp_mode eagle_with_att --mtp_step 2 + --linear_att_cache_size 500 + --pd_trans_mode nixl --pd_kv_page_size 4096 --pd_master_ip 127.0.0.1 --pd_master_port "$PORT" + --llm_prefill_att_backend fa3 triton --llm_decode_att_backend fa3 triton +) +P_ARGS=( + "${COMMON_ARGS[@]}" --run_mode prefill --tp "$PREFILL_TP" + --vit_att_backend triton + --graph_max_batch_size 8 --running_max_req_size 8 --batch_max_tokens 8192 + --chunked_prefill_size 8192 + --linear_att_hash_page_size 2048 --linear_att_page_block_num 8 + --enable_cpu_cache --cpu_cache_storage_size "${CPU_CACHE_SIZE:-2}" +) +D_ARGS=( + "${COMMON_ARGS[@]}" --run_mode decode --tp "$DECODE_TP" + --graph_max_batch_size "${GRAPH_MAX_BATCH:-64}" + --running_max_req_size "${MAX_REQUESTS:-64}" --batch_max_tokens 256 + --chunked_prefill_size 256 --port 28763 --nccl_port 29763 +) +MASTER_ARGS=( + "${MODEL_ARGS[@]}" --run_mode pd_master --pd_master_mode 2p1d --port "$PORT" +) + +cd -- "$LIGHTLLM_ROOT" +PIDS=() +cleanup() { + kill -TERM "${PIDS[@]}" 2>/dev/null || true + wait "${PIDS[@]}" 2>/dev/null || true +} +trap cleanup EXIT +trap 'exit 130' INT +trap 'exit 143' TERM + +start_server() { + local role="$1" gpu_ids="$2"; shift 2 + if [[ "$DRY_RUN" == 1 ]]; then + printf 'CUDA_VISIBLE_DEVICES=%q ' "$gpu_ids" + printf '%q ' python -m lightllm.server.api_server "$@" + printf '\n' + return + fi + printf 'Starting %s on GPUs [%s]\n' "$role" "$gpu_ids" + CUDA_VISIBLE_DEVICES="$gpu_ids" python -m lightllm.server.api_server "$@" & + PIDS+=("$!") +} +start_server prefill-1 "$P1_GPUS" "${P_ARGS[@]}" --port 28761 --nccl_port 29761 +start_server prefill-2 "$P2_GPUS" "${P_ARGS[@]}" --port 28762 --nccl_port 29762 +start_server decode "$D_GPUS" "${D_ARGS[@]}" +start_server pd-master '' "${MASTER_ARGS[@]}" +if [[ "$DRY_RUN" == 1 ]]; then + exit 0 +fi + +printf '2P1D is starting at http://127.0.0.1:%s\n' "$PORT" +wait -n "${PIDS[@]}" +exit 1 diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_mtp_diverse_autotune.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_mtp_diverse_autotune.py index 44664aee22..0989fe6d26 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_mtp_diverse_autotune.py +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_mtp_diverse_autotune.py @@ -183,7 +183,7 @@ def reference_attention(inputs): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize("kv_len", [17, 8192, 16384]) @pytest.mark.parametrize("block_batch", [1, 3, 4]) -@pytest.mark.parametrize("head_dim", [64, 128]) +@pytest.mark.parametrize("head_dim", [64, 128, 256]) def test_shared_groups_match_fp32_reference(monkeypatch, kv_len, block_batch, head_dim): monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", str(kv_len)) inputs = rebuild(make_inputs("cuda", block_batch=block_batch, head_dim=head_dim, token_count=7 * kv_len)) diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_stage2_autotune.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_stage2_autotune.py index c7f2e4babe..48956c4f36 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_stage2_autotune.py +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_stage2_autotune.py @@ -114,7 +114,7 @@ def reduction_reference(inputs): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize("block_n", [16, 64]) @pytest.mark.parametrize("capacity", [2, 128]) -@pytest.mark.parametrize("dtype,head_dim", [(torch.bfloat16, 128), (torch.float16, 64)]) +@pytest.mark.parametrize("dtype,head_dim", [(torch.bfloat16, 128), (torch.float16, 64), (torch.bfloat16, 256)]) def test_all_configs_reduce_only_valid_blocks(block_n, capacity, dtype, head_dim): lengths = [1, block_n, block_n + 1, capacity * block_n, capacity * block_n + 8192] mid = torch.randn(len(lengths), 4, capacity, head_dim, dtype=dtype, device="cuda") diff --git a/unit_tests/common/state_cache_manager/test_linear_att.py b/unit_tests/common/state_cache_manager/test_linear_att.py new file mode 100644 index 0000000000..bceb67f4f6 --- /dev/null +++ b/unit_tests/common/state_cache_manager/test_linear_att.py @@ -0,0 +1,64 @@ +import pytest +import torch + +from lightllm.common.state_cache_manager import LinearAttCacheConfig + + +@pytest.mark.parametrize("draft_layers", [0, 1, 3]) +@pytest.mark.parametrize( + "layer_types,expected_indices,linear_layers", + [ + ( + ("linear_attention", "linear_attention", "linear_attention", "full_attention") * 2, + (0, 1, 2, 0, 3, 4, 5, 1), + 6, + ), + ( + ( + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + ), + (0, 0, 1, 1, 2), + 3, + ), + ( + ("linear_attention", "full_attention", "linear_attention") * 2 + ("full_attention",), + (0, 0, 1, 2, 1, 3, 2), + 4, + ), + ], +) +def test_hybrid_layer_mapping_and_mtp_tail(layer_types, expected_indices, linear_layers, draft_layers): + config = LinearAttCacheConfig( + tp_world_size=1, + full_att_all_num_kv_heads=1, + full_att_dtype=torch.bfloat16, + full_att_num_kv_heads=1, + full_att_head_dim=128, + global_linear_k_heads=1, + global_linear_v_heads=1, + num_linear_k_heads=1, + num_linear_v_heads=1, + head_linear_k_dim=128, + head_linear_v_dim=128, + conv_kernel_size=4, + conv_state_dtype=torch.bfloat16, + ssm_state_dtype=torch.float32, + layer_types=layer_types, + draft_full_att_kv_layer_num=draft_layers, + ) + assert config.all_layer_num == len(layer_types) + assert config.linear_layer_num == linear_layers + full_layers = len(layer_types) - linear_layers + assert config.get_main_model_full_att_layer_num() == full_layers + assert config.get_full_att_kv_layer_num_with_draft_model() == full_layers + draft_layers + for layer, (kind, expected) in enumerate(zip(layer_types, expected_indices)): + get_index = ( + config.get_linear_state_layer_index if kind == "linear_attention" else config.get_full_att_kv_layer_index + ) + assert get_index(layer) == expected + for draft in range(draft_layers): + assert config.get_full_att_kv_layer_index(len(layer_types) + draft) == full_layers + draft diff --git a/unit_tests/models/glm5_next/test_cache.py b/unit_tests/models/glm5_next/test_cache.py index 2d923cb720..6ec706bd2a 100644 --- a/unit_tests/models/glm5_next/test_cache.py +++ b/unit_tests/models/glm5_next/test_cache.py @@ -67,11 +67,9 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy(monkeypatch, small_page, t head_linear_k_dim=128, head_linear_v_dim=128, conv_kernel_size=4, - linear_layer_num=3, conv_state_dtype=torch.bfloat16, ssm_state_dtype=torch.float32, - full_attention_interval=4, - all_layer_num=4, + layer_types=("linear_attention",) * 3 + ("deepseek_sparse_attention",), draft_full_att_kv_layer_num=int(mtp_step > 0), ) full_layers = config.get_full_att_kv_layer_num_with_draft_model() diff --git a/unit_tests/models/glm5_next/test_pd_cache.py b/unit_tests/models/glm5_next/test_pd_cache.py index 57665cbbe3..60f9bb7107 100644 --- a/unit_tests/models/glm5_next/test_pd_cache.py +++ b/unit_tests/models/glm5_next/test_pd_cache.py @@ -53,11 +53,9 @@ def create(tp): head_linear_k_dim=128, head_linear_v_dim=128, conv_kernel_size=4, - linear_layer_num=6, conv_state_dtype=torch.bfloat16, ssm_state_dtype=torch.float32, - full_attention_interval=4, - all_layer_num=8, + layer_types=(("linear_attention",) * 3 + ("deepseek_sparse_attention",)) * 2, draft_full_att_kv_layer_num=int(mtp_step > 0), ) mems = [] diff --git a/unit_tests/models/sensenova_loop/test_loop.py b/unit_tests/models/sensenova_loop/test_loop.py new file mode 100644 index 0000000000..9444485ab6 --- /dev/null +++ b/unit_tests/models/sensenova_loop/test_loop.py @@ -0,0 +1,278 @@ +import dataclasses +import json +from copy import deepcopy +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.models.registry import ModelRegistry +from lightllm.models.sensenova_loop.config import get_loop_layer_indices, get_loop_layer_types +from lightllm.common.state_cache_manager.linear_att import LinearAttCacheConfig +from lightllm.server.core.objs.start_args_type import StartArgs + + +def _config(layers=32, start=4, end=28, count=3): + return { + "model_type": "sensenova_loop", + "loop_settings": {"start": start, "end": end, "eval_count": count, "use_gate": False}, + "text_config": { + "num_hidden_layers": layers, + "full_attention_interval": 4, + "layer_types": ["full_attention" if (i + 1) % 4 == 0 else "linear_attention" for i in range(layers)], + "num_key_value_heads": 4, + "head_dim": 256, + "linear_num_key_heads": 16, + "linear_num_value_heads": 32, + "linear_key_head_dim": 128, + "linear_value_head_dim": 128, + "linear_conv_kernel_dim": 4, + }, + } + + +@pytest.mark.parametrize("start,end,count", [(4, 28, 3), (4, 28, 2), (1, 6, 2)]) +def test_cache_slots_follow_execution_depth_including_unaligned_loop(monkeypatch, start, end, count): + config = _config(start=start, end=end, count=count) + args = StartArgs(model_dir="unused", tp=4, dp=1, data_type="bfloat16", mtp_mode="eagle_with_att", mtp_step=3) + monkeypatch.setattr("lightllm.common.state_cache_manager.linear_att.get_env_start_args", lambda: args) + monkeypatch.setattr("lightllm.common.state_cache_manager.linear_att.get_added_mtp_kv_layer_num", lambda: 1) + monkeypatch.setattr("transformers.configuration_utils.PretrainedConfig.get_config_dict", lambda _: (config, {})) + cache = LinearAttCacheConfig.load_from_args() + indices = get_loop_layer_indices(config["loop_settings"], 32) + kinds = get_loop_layer_types(config["loop_settings"], config["text_config"]) + assert len(kinds) == 32 + (end - start) * (count - 1) + full = linear = 0 + slots = {} + for execution, (physical, kind) in enumerate(zip(indices, kinds)): + if kind == "full_attention": + slot = cache.get_full_att_kv_layer_index(execution) + assert slot == full + full += 1 + else: + slot = cache.get_linear_state_layer_index(execution) + assert slot == linear + linear += 1 + assert slot not in slots.setdefault(physical, set()) + slots[physical].add(slot) + assert cache.get_main_model_full_att_layer_num() == full + assert cache.linear_layer_num == linear + assert cache.get_full_att_kv_layer_index(len(indices)) == full # Independent MTP draft. + if (start, end, count) == (4, 28, 3): + assert (cache.all_layer_num, full, linear) == (80, 20, 60) + if (start, end, count) == (4, 28, 2): + assert (cache.all_layer_num, full, linear) == (56, 14, 42) + + +def test_physical_weights_are_created_and_loaded_once(monkeypatch): + from lightllm.models.sensenova_loop.model import SenseNovaLoopTpPartModel + from lightllm.models.qwen3_5.model import Qwen3_5TpPartModel + + model = object.__new__(SenseNovaLoopTpPartModel) + model.loop_layer_indices = get_loop_layer_indices(_config()["loop_settings"], 32) + model.num_physical_layers = 32 + model.config = {"n_layer": 80} + model.data_type, model.quant_cfg = torch.bfloat16, None + loaded = [] + + def create_weight(index, dtype, network_config, quant_cfg): + assert network_config is model.config + return object() + + def load_weights(self, weights): + loaded.extend(self.trans_layers_weight) + + model.pre_and_post_weight_class = lambda dtype, network_config: object() + model.transformer_weight_class = create_weight + monkeypatch.setattr(Qwen3_5TpPartModel, "load_weights", load_weights) + model._init_weights() + model.load_weights({}) + assert len(loaded) == len(set(loaded)) == 32 + assert len(model.trans_layers_weight) == 80 + for execution, physical in enumerate(model.loop_layer_indices): + assert model.trans_layers_weight[execution] is loaded[physical] + assert model.config["n_layer"] == 80 + + +@pytest.mark.parametrize("streams,count", [(2, 2), (4, 3)]) +def test_mhc_registry_and_embedded_mtp_reuse_existing_qwen_draft(tmp_path, streams, count): + from lightllm.models.draft_registry import get_draft_model_class + from lightllm.models.sensenova_loop.model import SenseNovaLoopTpPartModel + from lightllm.models.qwen3_5_mtp.model import Qwen3_5MTPModel + + config = _config(count=count) + config["mhc"] = dict(start=4, end=28, streams=streams, iterations=20) + config["text_config"].update(num_attention_heads=16, hidden_size=4096, vocab_size=128) + config["text_config"]["layer_types"] = ["linear_attention"] * 3 + ["full_attention"] + config["text_config"]["layer_types"] *= 8 + (tmp_path / "config.json").write_text(json.dumps(config)) + registration = ModelRegistry.get_model_config(config) + assert registration.model_class.endswith(":SenseNovaLoopTpPartModel") + assert registration.is_multimodal + assert get_draft_model_class(config, "eagle_with_att") is Qwen3_5MTPModel + + target = object.__new__(SenseNovaLoopTpPartModel) + target.weight_dir_, target.finetune_config, target.tp_world_size_ = str(tmp_path), None, 4 + target.args = SimpleNamespace(**dataclasses.asdict(StartArgs())) + target._init_config() + assert target.config["n_layer"] == 32 + 24 * (count - 1) + assert target.num_physical_layers == 32 + + draft = object.__new__(Qwen3_5MTPModel) + draft.weight_dir_, draft.finetune_config, draft.tp_world_size_ = str(tmp_path), None, 4 + draft._init_config() + assert draft.config["n_layer"] == draft.config["full_attention_interval"] == 1 + assert draft.config["layer_types"] == config["text_config"]["layer_types"] + + +def test_loop_tokenizer_loads_without_hf_loop_autoconfig(tmp_path): + from tokenizers import Tokenizer, models + from transformers import PreTrainedTokenizerFast, Qwen2VLImageProcessorFast + from lightllm.server.tokenizer import get_tokenizer + + tokenizer = PreTrainedTokenizerFast( + tokenizer_object=Tokenizer(models.WordLevel({"hello": 0, "[UNK]": 1}, unk_token="[UNK]")), + unk_token="[UNK]", + ) + tokenizer.save_pretrained(tmp_path) + Qwen2VLImageProcessorFast().save_pretrained(tmp_path) + config = _config() + config.update( + architectures=["SenseNovaLoopForConditionalGeneration"], + vision_config={}, + vision_start_token_id=2, + vision_end_token_id=3, + image_token_id=4, + ) + (tmp_path / "config.json").write_text(json.dumps(config)) + wrapped = get_tokenizer(str(tmp_path)) + assert wrapped.encode("hello") == [0] + assert isinstance(wrapped.image_processor, Qwen2VLImageProcessorFast) + + +@pytest.mark.parametrize("custom_layout", [False, True]) +def test_model_and_worker_use_the_same_loop_cache_layout(monkeypatch, custom_layout): + from lightllm.models.sensenova_loop.model import SenseNovaLoopTpPartModel + + config = _config(start=1, end=6, count=2) + if custom_layout: + config["text_config"]["layer_types"] = ["full_attention", "linear_attention"] * 16 + original = deepcopy(config) + args = StartArgs(tp=4, dp=1, data_type="bfloat16") + monkeypatch.setattr("lightllm.common.state_cache_manager.linear_att.get_added_mtp_kv_layer_num", lambda: 1) + monkeypatch.setattr("lightllm.common.state_cache_manager.linear_att.get_env_start_args", lambda: args) + monkeypatch.setattr("transformers.configuration_utils.PretrainedConfig.get_config_dict", lambda _: (config, {})) + monkeypatch.setattr( + "lightllm.models.qwen3next.model.ReqManagerForMamba", + lambda *args, linear_config: SimpleNamespace(linear_config=linear_config), + ) + raw_cache = LinearAttCacheConfig.load_from_args() + model = object.__new__(SenseNovaLoopTpPartModel) + model.args = args + model.max_req_num, model.batch_max_tokens, model.max_seq_length = 2, 8, 16 + model._init_req_manager() + assert dataclasses.asdict(raw_cache) == dataclasses.asdict(model.linear_config) + assert model.req_manager.linear_config is model.linear_config + assert config == original + + +def test_plain_loop_registry_does_not_enable_mhc(): + registration = ModelRegistry.get_model_config(_config()) + assert registration.model_class == "lightllm.models.sensenova_loop.model:SenseNovaLoopTpPartModel" + assert registration.is_multimodal + + +def test_weight_load_failure_restores_execution_weights(monkeypatch): + from lightllm.models.sensenova_loop.model import SenseNovaLoopTpPartModel + from lightllm.models.qwen3_5.model import Qwen3_5TpPartModel + + model = object.__new__(SenseNovaLoopTpPartModel) + model.config = {"n_layer": 80} + model.physical_layer_weights = [object()] + model.trans_layers_weight = [object()] + execution_config, execution_weights = model.config, model.trans_layers_weight + + def fail(*args): + raise RuntimeError("weight load failed") + + monkeypatch.setattr(Qwen3_5TpPartModel, "load_weights", fail) + with pytest.raises(RuntimeError, match="weight load failed"): + model.load_weights({}) + assert model.config is execution_config + assert model.trans_layers_weight is execution_weights + + +@pytest.fixture +def layer_runtime(monkeypatch): + from lightllm.utils.envs_utils import get_env_start_args + + monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(dataclasses.asdict(StartArgs()))) + for key, value in { + "DP_WORLD_SIZE": 1, + "CURRENT_RANK_IN_DP": 0, + "CURRENT_DEVICE_ID": 0, + "CURRENT_RANK_IN_NODE": 0, + }.items(): + monkeypatch.setenv("LIGHTLLM_" + key, str(value)) + get_env_start_args.cache_clear() + yield + get_env_start_args.cache_clear() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("mhc_enabled", [False, True]) +def test_loop_layer_types_are_local_to_sensenova(layer_runtime, tmp_path, mhc_enabled): + from lightllm.common.quantization import Quantcfg + from lightllm.models.qwen3_5.layer_infer.transformer_layer_infer import Qwen35TransformerLayerInfer + from lightllm.models.qwen3_5.layer_weights.transformer_layer_weight import Qwen35TransformerLayerWeight + from lightllm.models.qwen3_5_mtp.model import Qwen3_5MTPModel + from lightllm.models.sensenova_loop.layer_infer.transformer_layer_infer import SenseNovaLoopTransformerLayerInfer + from lightllm.models.sensenova_loop.layer_weights.transformer_layer_weight import ( + SenseNovaLoopTransformerLayerWeight, + ) + + config = _config(layers=8, start=1, end=6, count=2) + text = config["text_config"] + text.update( + n_layer=8, + hidden_size=128, + intermediate_size=256, + num_attention_heads=1, + num_key_value_heads=1, + head_dim=128, + linear_num_key_heads=1, + linear_num_value_heads=1, + rms_norm_eps=1e-6, + vocab_size=128, + layer_types=["full_attention", "linear_attention"] * 4, + ) + if mhc_enabled: + text["mhc"] = dict(start=1, end=6, streams=2, iterations=20) + indices = get_loop_layer_indices(config["loop_settings"], 8) + execution_config = { + **text, + "loop_physical_indices": indices, + "layer_types": get_loop_layer_types(config["loop_settings"], text), + } + weights = [ + SenseNovaLoopTransformerLayerWeight(i, torch.bfloat16, execution_config, Quantcfg(execution_config)) + for i in range(8) + ] + for execution, physical in enumerate(indices): + layer = SenseNovaLoopTransformerLayerInfer(execution, execution_config) + expected_linear = text["layer_types"][physical] == "linear_attention" + assert layer.is_linear_attention_layer == weights[physical].is_linear_attention_layer == expected_linear + assert hasattr(weights[physical], "linear_in_proj") == expected_linear + assert hasattr(weights[physical], "qkvo_gate_proj") != expected_linear + + # The base model still uses its original periodic layout, even with layer_types present. + base_layer = Qwen35TransformerLayerInfer(0, text) + base_weight = Qwen35TransformerLayerWeight(0, torch.bfloat16, text, Quantcfg(text)) + assert base_layer.is_linear_attention_layer and base_weight.is_linear_attention_layer + + (tmp_path / "config.json").write_text(json.dumps(config)) + draft = object.__new__(Qwen3_5MTPModel) + draft.weight_dir_, draft.finetune_config, draft.tp_world_size_ = str(tmp_path), None, 1 + draft._init_config() + for global_index in (len(indices), len(indices) + 1): + assert not Qwen35TransformerLayerInfer(global_index, draft.config).is_linear_attention_layer diff --git a/unit_tests/models/sensenova_loop/test_loop_mhc.py b/unit_tests/models/sensenova_loop/test_loop_mhc.py new file mode 100644 index 0000000000..34d4fe9696 --- /dev/null +++ b/unit_tests/models/sensenova_loop/test_loop_mhc.py @@ -0,0 +1,107 @@ +import pytest +import torch +import torch.nn.functional as F + +from lightllm.models.sensenova_loop.triton_kernel.mhc import hc_contract, hc_expand, hc_post, hc_pre_norm +from lightllm.models.sensenova_loop.triton_kernel.mhc.pre_norm import has_deepgemm_mhc +from lightllm.common.basemodel.triton_kernel.norm.rmsnorm import rmsnorm_forward + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +def _route_reference(x, weight, bias, streams, iterations): + # Training normalizes the residual before the FP32 router projection. + flat = x.float() + logits = F.linear(flat * torch.rsqrt(flat.square().mean(-1, keepdim=True) + 1e-6), weight, bias) + pre = logits[:, :streams].sigmoid() + post = 2 * logits[:, streams : 2 * streams].sigmoid() + logits = logits[:, 2 * streams :].reshape(-1, streams, streams) + for _ in range(iterations): + logits = logits - torch.logsumexp(logits, dim=-1, keepdim=True) + logits = logits - torch.logsumexp(logits, dim=-2, keepdim=True) + mix = logits.exp() + branch = (pre.to(x.dtype).unsqueeze(1) @ x.view(-1, streams, x.shape[1] // streams)).squeeze(1) + return branch, mix, post + + +@pytest.mark.parametrize("tokens", [1, 19, 257]) +@pytest.mark.parametrize("streams", [2, 4]) +def test_loop_mhc_matches_training_routing_and_rounding(tokens, streams): + torch.manual_seed(35) + hidden = 4096 + x = torch.randn(tokens, streams * hidden, device="cuda", dtype=torch.bfloat16) + weight = torch.randn(2 * streams + streams * streams, streams * hidden, device="cuda") * 0.008 + bias = torch.randn(weight.shape[0], device="cuda") + norm = torch.randn(hidden, device="cuda", dtype=torch.bfloat16) + expected = _route_reference(x, weight, bias, streams, 20) + actual = hc_pre_norm( + x, + weight, + bias, + norm, + streams, + 1e-6, + 1e-6, + 20, + use_deepgemm=False, + ) + torch.testing.assert_close(actual[0], rmsnorm_forward(expected[0], norm, 1e-6), atol=0.02, rtol=0.02) + torch.testing.assert_close(actual[1], expected[1], atol=2e-6, rtol=2e-5) + torch.testing.assert_close(actual[2], expected[2], atol=2e-6, rtol=2e-5) + layer_out = torch.randn(tokens, hidden, device="cuda", dtype=torch.bfloat16) + # Qwen uses output,input order; both matmul and the post product round to bf16. + expected_out = expected[1].to(x.dtype) @ x.view(tokens, streams, hidden) + expected_out += expected[2].to(x.dtype).unsqueeze(-1) * layer_out.unsqueeze(1) + actual_out = hc_post(layer_out, x, actual[1].transpose(1, 2), actual[2], streams) + torch.testing.assert_close(actual_out.view_as(expected_out), expected_out, atol=0.02, rtol=0.01) + + +@pytest.mark.skipif(not has_deepgemm_mhc(), reason="DeepGEMM mHC projection unavailable") +def test_loop_deepgemm_projection_matches_torch(): + x = torch.randn(19, 16384, device="cuda", dtype=torch.bfloat16) + weight = torch.randn(24, 16384, device="cuda") * 0.005 + bias = torch.randn(24, device="cuda") + norm = torch.randn(4096, device="cuda", dtype=torch.bfloat16) + actual = hc_pre_norm(x, weight, bias, norm, 4, 1e-6, 1e-6, 20) + expected = hc_pre_norm(x, weight, bias, norm, 4, 1e-6, 1e-6, 20, use_deepgemm=False) + for a, b in zip(actual, expected): + torch.testing.assert_close(a, b, atol=0.02, rtol=0.02) + + +@pytest.mark.parametrize("streams", [2, 4]) +def test_loop_mhc_decode_can_capture_and_replay_cuda_graph(streams): + x = torch.randn(3, streams * 4096, device="cuda", dtype=torch.bfloat16) + weight = torch.randn(2 * streams + streams * streams, streams * 4096, device="cuda") * 0.005 + bias = torch.randn(weight.shape[0], device="cuda") + norm = torch.ones(4096, device="cuda", dtype=torch.bfloat16) + + def forward(): + branch, mix, post = hc_pre_norm( + x, + weight, + bias, + norm, + streams, + 1e-6, + 1e-6, + 20, + use_deepgemm=False, + ) + return hc_post(branch, x, mix.transpose(1, 2), post, streams) + + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + forward() + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = forward() + x.normal_() + expected = forward() + graph.replay() + torch.testing.assert_close(output, expected, atol=0, rtol=0) + collapsed = torch.randn(3, 4096, device="cuda", dtype=torch.bfloat16) + assert torch.equal(hc_contract(hc_expand(collapsed, streams), streams), collapsed) diff --git a/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py b/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py index 7b60c15c4a..3111abce81 100644 --- a/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py +++ b/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py @@ -50,11 +50,9 @@ def linear_cache(monkeypatch, tmp_path, request): head_linear_k_dim=8, head_linear_v_dim=8, conv_kernel_size=3, - linear_layer_num=1, conv_state_dtype=torch.bfloat16, ssm_state_dtype=torch.float32, - full_attention_interval=2, - all_layer_num=2, + layer_types=("linear_attention", "full_attention"), ) mem = Qwen3NextMemManager(384, torch.bfloat16, 1, 8, 1, config) requests = ReqManagerForMamba(2, 512, mem, config) diff --git a/unit_tests/server/test_tokenizer.py b/unit_tests/server/test_tokenizer.py index 09d7ef8bc7..2accf0ce89 100644 --- a/unit_tests/server/test_tokenizer.py +++ b/unit_tests/server/test_tokenizer.py @@ -105,6 +105,7 @@ def load(*args, **kwargs): ({"model_type": "qwen3_vl", "vision_config": {}}, "qwen3_vl.model", "QWen3VLTokenizer"), ({"model_type": "qwen3_vl_moe", "vision_config": {}}, "qwen3_vl.model", "QWen3VLTokenizer"), ({"model_type": "qwen3_5", "vision_config": {}}, "qwen3_5.model", "QWen3_5Tokenizer"), + ({"model_type": "sensenova_loop", "vision_config": {}}, "qwen3_5.model", "QWen3_5Tokenizer"), ({"model_type": "qwen3_5_moe", "vision_config": {}}, "qwen3_5.model", "QWen3_5Tokenizer"), ({"thinker_config": {}}, "qwen3_omni_moe_thinker.model", "QWen3OmniTokenizer"), ({"model_type": "internvl_chat"}, "internvl.model", "InternvlTokenizer"), @@ -124,6 +125,7 @@ def test_selected_multimodal_tokenizer_is_loaded(monkeypatch, config, module_nam monkeypatch.setattr(module.AutoTokenizer, "from_pretrained", lambda *a, **kw: base_tokenizer) processor = SimpleNamespace(image_processor=object()) monkeypatch.setattr(transformers.AutoProcessor, "from_pretrained", lambda *a, **kw: processor) + monkeypatch.setattr(transformers.AutoImageProcessor, "from_pretrained", lambda *a, **kw: processor.image_processor) vision = ModuleType("lightllm.models.qwen2_vl.vision_process") vision.Qwen2VLImageProcessor = SimpleNamespace(from_pretrained=lambda *a, **kw: processor.image_processor) monkeypatch.setitem(sys.modules, vision.__name__, vision)