[Feature] support pooling model runner

fastdeploy/config.py

@@ -242,6 +242,9 @@ def _post_init(self):
 
         self.enable_mm = is_multimodal_model
 
+        if self.runner_type == "pooling":
+            os.environ["FD_USE_GET_SAVE_OUTPUT_V1"] = "1"
+
         if self.runner_type == "generate" and not is_generative_model:
             if is_multimodal_model:
                 pass

fastdeploy/engine/common_engine.py

@@ -737,6 +737,7 @@ def _fetch_request():
                                 raise
                 # 2. Schedule requests
                 tasks = self.resource_manager.schedule()
+
                 # 3. Send to engine
                 if tasks:
                     if self.cfg.scheduler_config.splitwise_role == "decode":

fastdeploy/model_executor/layers/pooler.py

@@ -305,19 +305,6 @@ def forward(self, pooled_data: Union[list[paddle.Tensor], paddle.Tensor], poolin
         return pooled_data
 
 
-def build_output(
-    all_data: Union[paddle.Tensor, list[paddle.Tensor]],
-) -> PoolerOutput:
-    # Pooling models D2H & synchronize occurs here
-    if isinstance(all_data, list):
-        all_data = [d.cpu() for d in all_data]
-    else:
-        all_data = all_data.cpu()
-
-    all_outputs = [PoolingSequenceGroupOutput(data) for data in all_data]
-    return PoolerOutput(outputs=all_outputs)
-
-
 class PoolingMethod(nn.Layer, ABC):
 
     @staticmethod
@@ -473,8 +460,11 @@ def forward(
         pooling_metadata: PoolingMetadata,
     ) -> PoolerOutput:
         pooled_data = self.extract_states(hidden_states, pooling_metadata)
-        pooled_data = self.head(pooled_data, pooling_metadata)
-        return build_output(pooled_data)
+        pooling_params = get_pooling_params(pooling_metadata)
+        assert len(pooled_data) == len(pooling_params)
+
+        pooled_data = [self.head(d, p) for d, p in zip(pooled_data, pooling_params)]
+        return pooled_data
 
 
 class SimplePooler(Pooler):
@@ -520,7 +510,7 @@ def forward(
     ) -> PoolerOutput:
         pooled_data = self.pooling(hidden_states, pooling_metadata)
         pooled_data = self.head(pooled_data, pooling_metadata)
-        return build_output(pooled_data)
+        return pooled_data
 
 
 class PoolerNormalize(PoolerActivation):
@@ -567,7 +557,7 @@ def forward(
                 hidden_states,
                 pooling_metadata[offset : offset + num_items],
             )
-            outputs.extend(group_output.outputs)
+            outputs.extend(group_output)
             offset += num_items
 
         return PoolerOutput(outputs)

fastdeploy/model_executor/pre_and_post_process.py

@@ -15,7 +15,7 @@
 """
 
 import queue
-from typing import Dict, Optional
+from typing import Dict, List, Optional, Union
 
 import numpy as np
 import paddle
@@ -85,6 +85,7 @@
         speculate_limit_thinking_content_length_v2,
     )
 
+from fastdeploy.output.pooler import PoolerOutput, PoolingSequenceGroupOutput
 from fastdeploy.output.stream_transfer_data import DecoderState, StreamTransferData
 from fastdeploy.worker.output import ModelOutputData, ModelRunnerOutput, SamplerOutput
 
@@ -238,17 +239,31 @@ def pre_process(
     )
 
 
-def _build_stream_transfer_data(output_tokens: np.ndarray):
+def _build_stream_transfer_data(output_tokens: np.ndarray, pooler_outputs: List[PoolingSequenceGroupOutput] = None):
     """Split output_tokens and output"""
-    output_tokens = output_tokens.reshape([-1]).numpy()
-    output_tokens_lists = np.split(output_tokens, output_tokens.shape[0])
 
     stream_transfer_datas = []
-    for bid, output_token_per_sample in enumerate(output_tokens_lists):
-        stream_transfer_data = StreamTransferData(
-            decoder_state=DecoderState.TEXT, tokens=output_token_per_sample, batch_id=bid
-        )
-        stream_transfer_datas.append(stream_transfer_data)
+    if output_tokens is not None:
+
+        output_tokens = output_tokens.reshape([-1]).numpy()
+        output_tokens_lists = np.split(output_tokens, output_tokens.shape[0])
+
+        for bid, output_token_per_sample in enumerate(output_tokens_lists):
+            stream_transfer_data = StreamTransferData(
+                decoder_state=DecoderState.TEXT, tokens=output_token_per_sample, batch_id=bid
+            )
+            stream_transfer_datas.append(stream_transfer_data)
+    elif pooler_outputs is not None:
+        for bid, pooler_output in enumerate(pooler_outputs):
+            if pooler_output.dtype == paddle.bfloat16:
+                pooler_output = pooler_output.astype("float32")
+
+            pooler_output = pooler_output.numpy()
+
+            stream_transfer_data = StreamTransferData(
+                decoder_state=DecoderState.TEXT, pooler_output=pooler_output, batch_id=bid
+            )
+            stream_transfer_datas.append(stream_transfer_data)
     return stream_transfer_datas
 
 
@@ -458,7 +473,7 @@ def post_process_specualate(
 
 
 def post_process(
-    sampler_output: SamplerOutput,
+    sampler_or_pooler_output: Union[SamplerOutput, PoolerOutput],
     model_output: ModelOutputData,
     share_inputs: Dict[str, paddle.Tensor],
     block_size: int = 64,
@@ -470,28 +485,40 @@ def post_process(
     line_break_id: int = -1,
 ) -> None:
     """Post-processing steps after completing a single token generation."""
-    if speculative_decoding:
-        post_process_specualate(
-            sampler_output,
-            model_output,
-            share_inputs,
-            save_each_rank,
-            skip_save_output,
-            think_end_id,
-            line_break_id,
-        )
-    else:
-        post_process_normal(
-            sampler_output,
+
+    if isinstance(sampler_or_pooler_output, PoolerOutput):
+        post_process_pooling(
+            sampler_or_pooler_output,
             model_output,
             share_inputs,
             block_size,
             save_each_rank,
             skip_save_output,
             async_output_queue,
-            think_end_id,
-            line_break_id,
         )
+    else:
+        if speculative_decoding:
+            post_process_specualate(
+                sampler_or_pooler_output,
+                model_output,
+                share_inputs,
+                save_each_rank,
+                skip_save_output,
+                think_end_id,
+                line_break_id,
+            )
+        else:
+            post_process_normal(
+                sampler_or_pooler_output,
+                model_output,
+                share_inputs,
+                block_size,
+                save_each_rank,
+                skip_save_output,
+                async_output_queue,
+                think_end_id,
+                line_break_id,
+            )
 
 
 def step_cuda(
@@ -775,3 +802,59 @@ def rebuild_padding(
     else:
         raise RuntimeError("Not supported platform")
     return hidden_states
+
+
+def post_process_pooling(
+    pooler_output: PoolerOutput,
+    model_output: ModelOutputData,
+    share_inputs: Dict[str, paddle.Tensor],
+    block_size: int = 64,
+    save_each_rank: bool = False,
+    skip_save_output: bool = False,
+    async_output_queue: queue.Queue = None,
+) -> None:
+
+    paddle.assign(
+        paddle.where(
+            model_output.stop_flags,
+            model_output.step_idx,
+            model_output.step_idx + 1,
+        ),
+        model_output.step_idx,
+    )
+    length_cond = paddle.greater_equal(model_output.step_idx, model_output.max_dec_len)
+
+    paddle.assign(
+        paddle.logical_or(model_output.stop_flags, length_cond),
+        model_output.stop_flags,
+    )
+
+    with paddle.framework._no_check_dy2st_diff():
+        if envs.ENABLE_V1_KVCACHE_SCHEDULER:
+            dummy_sampled_tokens = paddle.full_like(model_output.next_tokens, -1, dtype="int64")
+
+            paddle.assign(
+                paddle.ones_like(model_output.stop_flags, dtype="bool"),
+                model_output.stop_flags,
+            )
+            update_inputs_v1(
+                model_output.stop_flags,
+                model_output.not_need_stop,
+                model_output.seq_lens_this_time,
+                model_output.seq_lens_encoder,
+                model_output.seq_lens_decoder,
+                share_inputs["step_seq_lens_decoder"],
+                share_inputs["prompt_lens"],
+                dummy_sampled_tokens,
+                model_output.input_ids,
+                share_inputs["block_tables"],
+                model_output.stop_nums,
+                model_output.next_tokens,
+                model_output.is_block_step,
+                block_size,
+            )
+
+    if not skip_save_output:
+        if save_each_rank or model_output.mp_rank == 0:
+            output = _build_stream_transfer_data(output_tokens=None, pooler_outputs=pooler_output.outputs)
+            async_output_queue.put(output)

fastdeploy/output/stream_transfer_data.py

@@ -35,8 +35,8 @@ class StreamTransferData:
     """StreamTransferData"""
 
     decoder_state: DecoderState
-    tokens: np.array
     batch_id: int
+    tokens: Optional[np.array] = None
     speculaive_decoding: bool = False
     logprobs: Optional[np.array] = None
     accept_tokens: Optional[np.array] = None