feat(RL): add RL support for verl

docs/CN/source/tutorial/api_server_args.rst

@@ -294,6 +294,10 @@ PD 分离模式参数
 
     当输入图片超过该阈值时，LightLLM 会先自动将其缩放到该像素预算内，再继续后续流程。
 
+.. option:: --disable_image_resize
+
+    禁用对超过 ``--max_image_pixels`` 的图片的自动缩放。默认开启自动缩放。
+
 .. option:: --visual_infer_batch_size
 
     每次推理批次中处理的图像数量，默认为 ``1``
@@ -310,10 +314,6 @@ PD 分离模式参数
 
     ViT 的数据并行实例数量，默认为 ``1``
 
-.. option:: --visual_nccl_ports
-
-    为 ViT 构建分布式环境的 NCCL 端口列表，例如 29500 29501 29502，默认为 [29500]
-
 .. option:: --vit_att_backend
 
     设置 ViT 使用的注意力后端。可选值为：
@@ -496,9 +496,9 @@ PD 分离模式参数
     * ``triton``: 使用 torch 和 triton kernel（默认）
     * ``sglang_kernel``: 使用 sglang_kernel 实现
 
-.. option:: --return_all_prompt_logprobs
+.. option:: --enable_prompt_logprobs
 
-    返回所有提示 token 的 logprobs
+    启用 prompt top-k logprobs 捕获
 
 .. option:: --use_reward_model
 
@@ -549,14 +549,6 @@ DeepSeek 冗余专家参数
 监控和日志参数
 --------------
 
-.. option:: --disable_log_stats
-
-    禁用吞吐量统计日志记录
-
-.. option:: --log_stats_interval
-
-    记录统计信息的间隔（秒），默认为 ``10``
-
 .. option:: --health_monitor
 
     检查服务健康状态并在出错时重启

docs/EN/source/tutorial/api_server_args.rst

@@ -293,6 +293,10 @@ Multimodal Parameters
 
     If an input image exceeds this threshold, LightLLM automatically resizes it down to this pixel budget before continuing.
 
+.. option:: --disable_image_resize
+
+    Disable automatic resize for images exceeding ``--max_image_pixels``. Resize is enabled by default.
+
 .. option:: --visual_infer_batch_size
 
     Number of images processed in each inference batch, default is ``1``
@@ -309,10 +313,6 @@ Multimodal Parameters
 
     Number of data parallel instances for ViT, default is ``1``
 
-.. option:: --visual_nccl_ports
-
-    List of NCCL ports for ViT, e.g., 29500 29501 29502, default is [29500]
-
 .. option:: --vit_att_backend
 
     Set the attention backend for ViT. Available options:
@@ -497,9 +497,9 @@ Sampling and Generation Parameters
     * ``triton``: Use torch and triton kernel (default)
     * ``sglang_kernel``: Use sglang_kernel implementation
 
-.. option:: --return_all_prompt_logprobs
+.. option:: --enable_prompt_logprobs
 
-    Return logprobs for all prompt tokens
+    Enable prompt top-k logprobs capture
 
 .. option:: --use_reward_model
 
@@ -550,14 +550,6 @@ DeepSeek Redundant Expert Parameters
 Monitoring and Logging Parameters
 ---------------------------------
 
-.. option:: --disable_log_stats
-
-    Disable throughput statistics logging
-
-.. option:: --log_stats_interval
-
-    Interval for recording statistics (seconds), default is ``10``
-
 .. option:: --health_monitor
 
     Check service health status and restart on error

lightllm/common/basemodel/basemodel.py

@@ -36,6 +36,10 @@
 )
 from lightllm.common.triton_utils.autotuner import Autotuner
 from lightllm.utils.infer_utils import post_empty_cache
+from lightllm.utils.torch_memory_saver_utils import (
+    TorchMemorySaverWrapper,
+    MemoryTag,
+)
 from .attention import get_prefill_att_backend_class, get_decode_att_backend_class
 from .attention import BaseAttBackend
 
@@ -93,6 +97,7 @@ def __init__(self, kvargs):
         self.tp_world_size_ = get_dp_world_size()
         self.enable_tpsp_mix_mode = get_env_start_args().enable_tpsp_mix_mode
 
+        self.torch_memory_saver = TorchMemorySaverWrapper(self.args.enable_torch_memory_saver)
         self.is_mtp_mode = self.args.mtp_mode in [
             "vanilla_with_att",
             "eagle_with_att",
@@ -106,18 +111,21 @@ def __init__(self, kvargs):
         self._verify_params()
         self._init_quant()
 
-        self._init_weights()
-        self._init_req_manager()
-        self._init_mem_manager()
+        enable_weight_cpu_backup = self.args.enable_weight_cpu_backup
+        with self.torch_memory_saver.region(tag=MemoryTag.WEIGHT, enable_cpu_backup=enable_weight_cpu_backup):
+            self._init_weights()
+        with self.torch_memory_saver.region(tag=MemoryTag.KV_CACHE):
+            self._init_req_manager()
+            self._init_mem_manager()
+
         # 因为类似 qwen3.5 的linear 架构的模型，其 req_manager 会存储运行时使用的大量 linear state
         # 这可能会占用大量的显存，所以，req_manger 中保存的 mem_manger 是mem manager 初始化后再赋值
         self.req_manager.mem_manager = self.mem_manager
-
         self._check_mem_size()
         self._init_infer_layer()
         self._init_some_value()
         self._init_custom()
-        self._load_hf_weights()
+        self.load_weights(self.weight_dict)
 
         self._init_att_backend()
         self._init_att_backend1()
@@ -176,13 +184,14 @@ def _init_weights(self, start_layer_index=0):
         ]
         return
 
-    def _load_hf_weights(self):
+    def load_weights(self, weight_dict: dict):
+        assert weight_dict is None or isinstance(weight_dict, dict), "weight_dict must be a dict or None"
         load_hf_weights(
-            self.data_type,
+            data_type=self.data_type,
             weight_dir=self.weight_dir_,
             pre_post_layer=self.pre_post_weight,
             transformer_layer_list=self.trans_layers_weight,
-            weight_dict=self.weight_dict,
+            weight_dict=weight_dict,
         )
         self.pre_post_weight.verify_load()
         [weight.verify_load() for weight in self.trans_layers_weight]
@@ -520,13 +529,13 @@ def _create_unpad_decode_model_output(self, model_output: ModelOutput, origin_ba
     def _create_unpad_prefill_model_output(
         self, padded_model_output: ModelOutput, origin_handle_token_num: int, origin_batch_size: int
     ):
-        if self.return_all_prompt_logics:
-            new_model_output = copy.copy(padded_model_output)
-            new_model_output.logits = new_model_output.logits[0:origin_handle_token_num]
-        else:
-            new_model_output = copy.copy(padded_model_output)
-            # 移除多余的pad 的那个 req 对应的 logics
-            new_model_output.logits = new_model_output.logits[0:origin_batch_size]
+        new_model_output = copy.copy(padded_model_output)
+        # logits 始终只对应每个请求最后一个位置，移除 padding 的 req 对应的行。
+        new_model_output.logits = new_model_output.logits[0:origin_batch_size]
+        # prompt_logics 保存整个 prefill 阶段所有 token 位置的 logits，
+        # 按实际处理的 token 数量裁剪掉 padding 部分（仅 return_all_prompt_logics 模式下非空）。
+        if new_model_output.prompt_logics is not None:
+            new_model_output.prompt_logics = new_model_output.prompt_logics[0:origin_handle_token_num]
 
         # 特殊模型，特殊模式的特殊变量的特殊 unpad
         if new_model_output.mtp_main_output_hiddens is not None:
@@ -701,7 +710,7 @@ 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)
-        model_output = ModelOutput(logits=predict_logits)
+        model_output = ModelOutput(logits=predict_logits, prompt_logics=infer_state.prompt_logics)
 
         # 特殊模型特殊模式的额外输出
         if self.is_mtp_mode:
@@ -974,8 +983,8 @@ def _overlap_tpsp_context_forward(self, infer_state: InferStateInfo, infer_state
         )
         g_cache_manager.cache_env_out()
 
-        model_output = ModelOutput(logits=predict_logits.contiguous())
-        model_output1 = ModelOutput(logits=predict_logits1.contiguous())
+        model_output = ModelOutput(logits=predict_logits.contiguous(), prompt_logics=infer_state.prompt_logics)
+        model_output1 = ModelOutput(logits=predict_logits1.contiguous(), prompt_logics=infer_state1.prompt_logics)
 
         if self.is_mtp_mode:
             input_embs = self.pre_infer._tpsp_allgather(input=input_embs, infer_state=infer_state)
@@ -1087,6 +1096,7 @@ def _check_max_len_infer(self):
             )
             logger.error(exception_str)
             raise Exception(exception_str)
+        torch.cuda.empty_cache()
         return
 
     def autotune_layers(self):
@@ -1221,6 +1231,9 @@ def _init_padded_req(self):
         del b_seq_len
         del b_ready_cache_len
         del model_output
+        del b_mtp_index
+        del b_prefill_start_loc
+        del b_q_seq_len
         torch.cuda.empty_cache()
         return
 

lightllm/common/basemodel/batch_objs.py

@@ -112,6 +112,12 @@ class ModelOutput:
     # 输入
     mtp_main_output_hiddens: Optional[torch.Tensor] = None
 
+    # prompt_logics 用于在开启 return_all_prompt_logics 模式（如 enable_prompt_logprobs）时，
+    # 保存整个 prefill 阶段每一个 token 位置对应的 logits（而非仅最后一个位置的 logits）。
+    # 此时 logits 依然只保存每个请求最后一个位置的 logits，prompt_logics 为可选项，仅在
+    # 需要返回 prompt logprobs 信息时才会非空。
+    prompt_logics: Optional[torch.Tensor] = None
+
     def to_no_ref_tensor(self):
         self.logits = tensor_to_no_ref_tensor(self.logits)
         if self.mtp_main_output_hiddens is not None:

lightllm/common/basemodel/cuda_graph.py

@@ -8,6 +8,10 @@
 from lightllm.utils.envs_utils import get_env_start_args
 from lightllm.distributed import dist_group_manager
 from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput
+from lightllm.utils.torch_memory_saver_utils import (
+    TorchMemorySaverWrapper,
+    MemoryTag,
+)
 from .infer_struct import InferStateInfo
 
 
@@ -54,6 +58,7 @@ def __init__(self, max_batch_size=8, max_len_in_batch=8192, tp_world_size: int =
         self.max_batch_size = max_batch_size
         self.graph_max_len_in_batch = max_len_in_batch
         self.enable_decode_microbatch_overlap = self.args.enable_decode_microbatch_overlap
+        self.torch_memory_saver = TorchMemorySaverWrapper(self.args.enable_torch_memory_saver)
 
         self.cuda_graph_batch_sizes = self.gen_cuda_graph_batch_sizes(
             max_batch_size=max_batch_size,
@@ -105,7 +110,7 @@ def _capture_decode(self, decode_func, infer_state: InferStateInfo):
                 if param_name not in pure_para_set:
                     delattr(infer_state, param_name)
 
-        with torch.cuda.graph(graph_obj, pool=self.mempool):
+        with self.torch_memory_saver.cuda_graph(graph_obj, pool=self.mempool):
             model_output = decode_func(infer_state)
         self.graph[batch_size] = (graph_obj, infer_state, model_output)
         graph_obj.replay()
@@ -139,7 +144,7 @@ def _capture_decode_overlap(
                 if para_name not in pure_para_set1:
                     delattr(infer_state1, para_name)
 
-        with torch.cuda.graph(graph_obj, pool=self.mempool):
+        with self.torch_memory_saver.cuda_graph(graph_obj, pool=self.mempool):
             model_output, model_output1 = decode_func(infer_state, infer_state1)
         self.graph[batch_size] = (
             graph_obj,

lightllm/common/basemodel/infer_struct.py

@@ -56,6 +56,10 @@ def __init__(self):
 
         self.is_token_healing: bool = False
         self.return_all_prompt_logics: bool = False
+        # 在开启 return_all_prompt_logics 模式时，保存整个 prefill 阶段每一个
+        # token 位置的 logits，供后续回传 prompt logprobs 信息使用。
+        # 仅在 prefill 阶段且需要返回 prompt logprobs 时才会被填充。
+        self.prompt_logics: Optional[torch.Tensor] = None
         self.multimodal_params: dict = None
         self.is_cuda_graph: bool = False  # 标记是否是cuda graph的捕获推理
         self.dist_group: CustomProcessGroup = None

lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py

@@ -9,6 +9,7 @@
     SliceMixinTpl,
 )
 from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.impl import select_fuse_moe_impl
+from lightllm.common.basemodel.moe_route_info_manager import get_moe_capture_callback
 from lightllm.common.quantization.quantize_method import QuantizationMethod
 from lightllm.utils.envs_utils import get_redundancy_expert_ids, get_redundancy_expert_num, get_env_start_args
 from lightllm.utils.dist_utils import get_global_world_size, get_global_rank
@@ -134,9 +135,11 @@ def experts(
         topk_group: int,
         num_expert_group: int,
         is_prefill: Optional[bool] = None,
+        infer_state=None,
         shared_expert_gate: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
-        """Backward compatible method that routes to platform-specific implementation."""
+        # Captures MoE topk expert ids for routed-experts metadata when enabled.
+        moe_capture_callback = get_moe_capture_callback(infer_state, self.layer_num_)
         return self.fuse_moe_impl(
             input_tensor=input_tensor,
             router_logits=router_logits,
@@ -150,6 +153,7 @@ def experts(
             topk_group=topk_group,
             num_expert_group=num_expert_group,
             is_prefill=is_prefill,
+            moe_capture_callback=moe_capture_callback,
             per_expert_scale=self.per_expert_scale,
             shared_expert_gate=shared_expert_gate,
         )
@@ -319,6 +323,7 @@ def _create_weight(self):
             device_id=self.device_id_,
             num_experts=self.local_n_routed_experts,
         )
+        self.w1, self.w3 = w13_param_list
         self.w1_list: List[WeightPack] = self._get_expert_weight_list(w13_param_list[0])
         self.w3_list: List[WeightPack] = self._get_expert_weight_list(w13_param_list[1])
         self.w2_list: List[WeightPack] = self._get_expert_weight_list(self.w2)
@@ -341,7 +346,6 @@ def _get_expert_weight_list(self, weight_pack: WeightPack):
         return weight_list
 
     def _load_weight(self, expert_idx_to_local_idx: Dict[int, int], weights: Dict[str, torch.Tensor]):
-        # Load each expert with TP slicing
         for expert_idx, local_expert_idx in expert_idx_to_local_idx.items():
             with self.lock:
                 self._load_expert(expert_idx, local_expert_idx, weights)

lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/gpt_oss_fused_moe_weight_tp.py

@@ -4,6 +4,7 @@
 from typing import Optional, Tuple, List, Dict, Any
 
 from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.fused_moe_weight import FusedMoeWeight
+from lightllm.common.basemodel.moe_route_info_manager import get_moe_capture_callback
 from lightllm.utils.dist_utils import get_current_rank_in_dp, get_current_device_id
 from lightllm.common.quantization import Quantcfg
 from lightllm.common.quantization.quantize_method import QuantizationMethod
@@ -144,12 +145,18 @@ def experts(
         topk_group: int,
         num_expert_group: int,
         is_prefill: Optional[bool] = None,
+        infer_state=None,
         shared_expert_gate: Optional[torch.Tensor] = None,
     ):
         assert shared_expert_gate is None, "shared_expert_gate is not supported by GPT-OSS fused MoE"
 
         topk_weights, topk_ids = self._router(router_logits, top_k)
 
+        # Captures MoE topk expert ids for routed-experts metadata when enabled.
+        moe_capture_callback = get_moe_capture_callback(infer_state, self.layer_num_)
+        if moe_capture_callback is not None:
+            moe_capture_callback(topk_ids)
+
         w1, w1_scale = self.w1
         w2, w2_scale = self.w2
         use_fp8_w8a8 = self.quant_method is not None

lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/base_impl.py

@@ -1,10 +1,10 @@
 import torch
 from abc import abstractmethod
+from typing import Callable, Optional
 from lightllm.common.quantization.quantize_method import (
     WeightPack,
     QuantizationMethod,
 )
-from typing import Optional
 from lightllm.utils.dist_utils import (
     get_global_rank,
     get_global_world_size,
@@ -62,6 +62,8 @@ def __call__(
         topk_group: int,
         num_expert_group: int,
         is_prefill: Optional[bool] = None,
+        # Callback to capture MoE topk expert ids (routed experts metadata).
+        moe_capture_callback: Optional[Callable[[torch.Tensor], None]] = None,
         per_expert_scale: Optional[torch.Tensor] = None,
         # Qwen3.5 uses this gate to control fused shared expert aggregation weights.
         shared_expert_gate: Optional[torch.Tensor] = None,