add logits compute into FD

Author: liuruyan

fastdeploy/model_executor/layers/sample/logprobs.py

@@ -133,6 +133,7 @@ def build_output_logprobs(
     is_naive: bool = False,
     logprobs_mode: str = "default",
     compute_logprobs_fn: Optional[Callable] = None,
+    compute_logits_stats: bool = False,
 ) -> Tuple[Optional[LogprobsTensors], Optional[paddle.Tensor]]:
     """
     Build logprobs output for both NAIVE and speculative (MTP/Ngram) modes.
@@ -151,6 +152,7 @@ def build_output_logprobs(
         logprobs_mode: One of "raw_logprobs", "raw_logits", or "default".
         compute_logprobs_fn: Callable for computing logprobs with temperature
             scaling and top_p normalization. Used when logprobs_mode == "raw_logprobs".
+        compute_logits_stats: Whether to compute per-token logits statistics (min/max/mean/std).
 
     Returns:
         tuple: (logprobs_tensors, cu_batch_token_offset)
@@ -218,4 +220,21 @@ def build_output_logprobs(
 
     logprobs_tensors = gather_logprobs(raw_logprobs, num_logprobs, token_ids=token_ids)
 
+    # Compute logits statistics (min/max/mean/std) per token and pack into LogprobsTensors
+    if compute_logits_stats:
+        with paddle.no_grad():
+            logits_min = paddle.min(output_logits, axis=1)
+            logits_max = paddle.max(output_logits, axis=1)
+            logits_mean = paddle.mean(output_logits, axis=1)
+            logits_std = paddle.std(output_logits, axis=1)
+        logprobs_tensors = LogprobsTensors(
+            logprob_token_ids=logprobs_tensors.logprob_token_ids,
+            logprobs=logprobs_tensors.logprobs,
+            selected_token_ranks=logprobs_tensors.selected_token_ranks,
+            logits_min=logits_min,
+            logits_max=logits_max,
+            logits_mean=logits_mean,
+            logits_std=logits_std,
+        )
+
     return logprobs_tensors, cu_batch_token_offset

fastdeploy/model_executor/layers/sample/sampler.py

@@ -375,6 +375,7 @@ def __init__(self, fd_config: FDConfig = None, logprobs_mode: str = "raw_logprob
 
         self.guided_decoding = GuidedDecoding(fd_config)
         self.logprobs_mode = fd_config.model_config.logprobs_mode if fd_config is not None else logprobs_mode
+        self.compute_logits_stats = fd_config.model_config.compute_logits_stats if fd_config is not None else False
         # Can only be created when fd_config.early_stopper_config.enable_early_stop = True
         if (
             fd_config is not None
@@ -522,6 +523,19 @@ def forward_cuda(
             elif self.logprobs_mode == "raw_logits":
                 raw_logprobs = logits.clone()
 
+        # Compute logits statistics (min/max/mean/std) per sequence before penalties
+        logits_min = None
+        logits_max = None
+        logits_mean = None
+        logits_std = None
+        if num_logprobs is not None and self.compute_logits_stats:
+            with paddle.no_grad():
+                # logits shape: [batch_size, vocab_size], compute stats per sequence (reduce over vocab dimension)
+                logits_min = paddle.min(logits, axis=1)  # [batch_size]
+                logits_max = paddle.max(logits, axis=1)  # [batch_size]
+                logits_mean = paddle.mean(logits, axis=1)  # [batch_size]
+                logits_std = paddle.std(logits, axis=1)  # [batch_size]
+
         for proc in sampling_metadata.logits_processors or []:
             logits = proc.apply(logits)
 
@@ -565,6 +579,33 @@ def forward_cuda(
         logprobs_tensors = (
             None if num_logprobs is None else self.gather_logprobs(raw_logprobs, num_logprobs, token_ids=next_tokens)
         )
+
+        # Pack logits stats into LogprobsTensors
+        if logprobs_tensors is not None and logits_min is not None:
+            if current_platform.is_cuda():
+                logits_min_cpu = paddle.empty_like(logits_min, device="cpu").pin_memory()
+                logits_max_cpu = paddle.empty_like(logits_max, device="cpu").pin_memory()
+                logits_mean_cpu = paddle.empty_like(logits_mean, device="cpu").pin_memory()
+                logits_std_cpu = paddle.empty_like(logits_std, device="cpu").pin_memory()
+                logits_min_cpu.copy_(logits_min, False)
+                logits_max_cpu.copy_(logits_max, False)
+                logits_mean_cpu.copy_(logits_mean, False)
+                logits_std_cpu.copy_(logits_std, False)
+            else:
+                logits_min_cpu = logits_min.cpu()
+                logits_max_cpu = logits_max.cpu()
+                logits_mean_cpu = logits_mean.cpu()
+                logits_std_cpu = logits_std.cpu()
+            logprobs_tensors = LogprobsTensors(
+                logprob_token_ids=logprobs_tensors.logprob_token_ids,
+                logprobs=logprobs_tensors.logprobs,
+                selected_token_ranks=logprobs_tensors.selected_token_ranks,
+                logits_min=logits_min_cpu,
+                logits_max=logits_max_cpu,
+                logits_mean=logits_mean_cpu,
+                logits_std=logits_std_cpu,
+            )
+
         if sampling_metadata.enable_early_stop:
             # will set the stop batch in stop_flags
             assert sampling_metadata.stop_flags is not None, "need stop_flags for early stop"
@@ -640,6 +681,7 @@ def __init__(self, fd_config: FDConfig):
         else:
             raise NotImplementedError
         self.logprobs_mode = fd_config.model_config.logprobs_mode
+        self.compute_logits_stats = fd_config.model_config.compute_logits_stats
         self.speculative_verify_window = fd_config.speculative_config.verify_window
         self.speculative_max_candidate_len = fd_config.speculative_config.max_candidate_len
         self.speculative_benchmark_mode = fd_config.speculative_config.benchmark_mode
@@ -1038,6 +1080,7 @@ def forward_cuda(
                 is_naive=is_naive,
                 logprobs_mode=self.logprobs_mode,
                 compute_logprobs_fn=self.compute_logprobs,
+                compute_logits_stats=self.compute_logits_stats,
             )
             sampler_output.logprobs_tensors = logprobs_tensors
             if cu_batch_token_offset is not None:
@@ -1147,6 +1190,7 @@ def __init__(self, fd_config: FDConfig):
         else:
             raise NotImplementedError
         self.logprobs_mode = fd_config.model_config.logprobs_mode
+        self.compute_logits_stats = fd_config.model_config.compute_logits_stats
         self.enable_draft_logprob = fd_config.speculative_config.enable_draft_logprob
 
     def pre_process(self, skip_idx_list: List[int] = []):
@@ -1336,6 +1380,24 @@ def forward_cuda(
 
             logprobs_tensors = self.gather_logprobs(raw_logprobs, num_logprobs, token_ids=token_ids)
 
+            # Compute logits statistics on draft_logits for MTP tokens
+            if self.compute_logits_stats:
+                draft_logits_for_stats = share_inputs["draft_logits"][:real_token_num, :]
+                with paddle.no_grad():
+                    logits_min = paddle.min(draft_logits_for_stats, axis=1)
+                    logits_max = paddle.max(draft_logits_for_stats, axis=1)
+                    logits_mean = paddle.mean(draft_logits_for_stats, axis=1)
+                    logits_std = paddle.std(draft_logits_for_stats, axis=1)
+                logprobs_tensors = LogprobsTensors(
+                    logprob_token_ids=logprobs_tensors.logprob_token_ids,
+                    logprobs=logprobs_tensors.logprobs,
+                    selected_token_ranks=logprobs_tensors.selected_token_ranks,
+                    logits_min=logits_min,
+                    logits_max=logits_max,
+                    logits_mean=logits_mean,
+                    logits_std=logits_std,
+                )
+
         sampler_output = SamplerOutput(
             sampled_token_ids=token_ids,
             logprobs_tensors=logprobs_tensors,

fastdeploy/worker/input_batch.py

@@ -21,6 +21,7 @@
 from fastdeploy.model_executor.layers.rotary_embedding import get_rope
 from fastdeploy.model_executor.logits_processor import build_logits_processors
 from fastdeploy.platforms import current_platform
+from fastdeploy.worker.output import LogprobsTensors
 
 
 class InputBatch:
@@ -1150,9 +1151,26 @@ def recover_batch_index_for_sampler_output(sampler_output, index_to_batch_id, en
         real_logprob_token_ids = _recover_tensor(logprob_token_ids, src_order)
         real_logprobs = _recover_tensor(logprobs, src_order)
         real_selected_token_ranks = _recover_tensor(selected_token_ranks, src_order)
-        sampler_output.logprobs_tensors.logprob_token_ids = real_logprob_token_ids
-        sampler_output.logprobs_tensors.logprobs = real_logprobs
-        sampler_output.logprobs_tensors.sampled_token_ranks = real_selected_token_ranks
+
+        real_logits_min = None
+        real_logits_max = None
+        real_logits_mean = None
+        real_logits_std = None
+        if sampler_output.logprobs_tensors.logits_min is not None:
+            real_logits_min = _recover_tensor(sampler_output.logprobs_tensors.logits_min, src_order)
+            real_logits_max = _recover_tensor(sampler_output.logprobs_tensors.logits_max, src_order)
+            real_logits_mean = _recover_tensor(sampler_output.logprobs_tensors.logits_mean, src_order)
+            real_logits_std = _recover_tensor(sampler_output.logprobs_tensors.logits_std, src_order)
+
+        sampler_output.logprobs_tensors = LogprobsTensors(
+            logprob_token_ids=real_logprob_token_ids,
+            logprobs=real_logprobs,
+            selected_token_ranks=real_selected_token_ranks,
+            logits_min=real_logits_min,
+            logits_max=real_logits_max,
+            logits_mean=real_logits_mean,
+            logits_std=real_logits_std,
+        )
 
     if sampler_output.token_num_per_batch is not None:
         token_num_per_batch = sampler_output.token_num_per_batch