[TRTLLM-12982][feat] support multi item scoring in LLM.encode (#14693)

Signed-off-by: ixlmar <206748156+ixlmar@users.noreply.github.com>

Author: ixlmar

tensorrt_llm/_torch/attention_backend/flashinfer.py

@@ -748,6 +748,13 @@ class AttentionForwardArgs:
     attention_mask_data: Optional[torch.Tensor] = None
     attention_sinks: Optional[torch.Tensor] = None
 
+    multi_item_part_lens: Optional[list[list[int]]] = None
+    """Additional token layout information for multi-item scoring.
+
+    Aggregates `TokensPrompt.multi_item_part_lens` for all requests in the batch,
+    see `TokensPrompt` for details.
+    """
+
     latent_cache: Optional[torch.Tensor] = None
     q_pe: Optional[torch.Tensor] = None
     mrope_rotary_cos_sin: Optional[torch.Tensor] = None

tensorrt_llm/_torch/attention_backend/interface.py

@@ -322,6 +322,10 @@ def forward(self,
         )
         assert not metadata.is_cross, "Star Attention does not support cross attention yet."
 
+        if forward_args.multi_item_part_lens is not None:
+            raise ValueError(
+                "Star Attention does not support multi-item scoring")
+
         q = q.view(-1, self.num_heads, self.head_dim)
         k = k.view(-1, self.num_kv_heads, self.head_dim)
         v = v.view(-1, self.num_kv_heads, self.head_dim)

tensorrt_llm/_torch/attention_backend/star_flashinfer.py

@@ -1718,6 +1718,10 @@ def forward(
         )
         assert not metadata.is_cross, "TRT-LLM Attention does not support cross attention yet."
 
+        if forward_args.multi_item_part_lens is not None:
+            raise ValueError(
+                "TRT-LLM Attention does not support multi-item scoring")
+
         # SM90 forces ``use_paged_context_fmha`` on for correctness
         # (https://nvbugs/5624818).
         if get_sm_version() == 90:

tensorrt_llm/_torch/attention_backend/trtllm.py

@@ -382,6 +382,11 @@ def forward(self,
                 forward_args: Optional[AttentionForwardArgs] = None,
                 **kwargs) -> torch.Tensor:
         forward_args = merge_attention_forward_args(forward_args, kwargs)
+
+        if forward_args.multi_item_part_lens is not None:
+            raise ValueError(
+                "Vanilla Attention does not support multi-item scoring")
+
         if metadata.kv_cache_manager is None:
             # NOTE: WAR for no kv cache attn e.g. BERT,
             # try to separate the kv cache estimation path from no kv cache attn.

tensorrt_llm/_torch/attention_backend/vanilla.py

@@ -699,6 +699,7 @@ def _attn_impl(
         output_sf: Optional[torch.Tensor] = None,
         attention_sinks: Optional[torch.Tensor] = None,
         has_lora: bool = False,
+        multi_item_part_lens: Optional[list[list[int]]] = None,
     ):
         num_tokens = attn_metadata.num_tokens
 
@@ -735,6 +736,7 @@ def _attn_impl(
                     attention_mask_data=attention_mask_data,
                     softmax_stats_tensor=softmax_stats,
                     attention_sinks=attention_sinks,
+                    multi_item_part_lens=multi_item_part_lens,
                 ))
             if isinstance(attn_output, tuple):
                 attn_output = attn_output[0]
@@ -781,6 +783,7 @@ def _attn_impl(
                 output=output[:num_tokens, :] if output is not None else None,
                 output_sf=output_sf,
                 attention_sinks=attention_sinks,
+                multi_item_part_lens=multi_item_part_lens,
             ))
         if isinstance(attn_output, tuple):
             assert len(
@@ -801,6 +804,7 @@ def forward_impl(
         mrope_config: Optional[dict],
         attention_sinks: Optional[torch.Tensor] = None,
         has_lora: bool = False,
+        multi_item_part_lens: Optional[list[list[int]]] = None,
     ):
         mrope_rotary_cos_sin = None
         mrope_position_deltas = None
@@ -837,17 +841,20 @@ def forward_impl(
                 output_sf,
             )
         else:
-            output, output_sf = self._attn_impl(q,
-                                                k,
-                                                v,
-                                                attn_metadata,
-                                                attention_mask,
-                                                mrope_rotary_cos_sin,
-                                                mrope_position_deltas,
-                                                attention_window_size,
-                                                attention_mask_data,
-                                                attention_sinks=attention_sinks,
-                                                has_lora=has_lora)
+            output, output_sf = self._attn_impl(
+                q,
+                k,
+                v,
+                attn_metadata,
+                attention_mask,
+                mrope_rotary_cos_sin,
+                mrope_position_deltas,
+                attention_window_size,
+                attention_mask_data,
+                attention_sinks=attention_sinks,
+                has_lora=has_lora,
+                multi_item_part_lens=multi_item_part_lens,
+            )
         if output_sf is not None:
             output = Fp4QuantizedTensor(output, output_sf)
 
@@ -865,6 +872,7 @@ def forward(
         attention_window_size: Optional[int] = None,
         attention_mask_data: Optional[torch.Tensor] = None,
         attention_sinks: Optional[torch.Tensor] = None,
+        multi_item_part_lens: Optional[list[list[int]]] = None,
         **kwargs,
     ) -> torch.Tensor:
         """
@@ -924,6 +932,23 @@ def forward(
             position_ids = self._adjust_position_ids_for_spec_dec(
                 position_ids, attn_metadata)
 
+        if multi_item_part_lens is not None:
+            # adjust RoPE positions for multi-item scoring
+            current_idx = 0
+            for req_multi_item_part_lens in multi_item_part_lens:
+                req_prefix_len, *req_multi_item_part_lens = req_multi_item_part_lens
+                # RoPE for prefix does not need updating and RoPE for delimiter does not matter
+                current_idx += req_prefix_len + 1
+                for item_len in req_multi_item_part_lens:
+                    next_idx = current_idx + item_len
+                    position_ids[0, current_idx:next_idx].copy_(
+                        torch.arange(req_prefix_len,
+                                     req_prefix_len + item_len,
+                                     dtype=position_ids.dtype,
+                                     device=position_ids.device),
+                        non_blocking=True)
+                    current_idx = next_idx + 1  # RoPE for delimiter does not matter
+
         q, k, v = self.apply_rope(q, k, v, position_ids)
         q, k, v = self.convert_qkv(q, k, v)
 
@@ -932,16 +957,19 @@ def forward(
                 f"Attention sinks are only supported with attn_backend='TRTLLM'. "
                 f"Current backend: {self.attn_backend}.")
 
-        attn_output = self.forward_impl(q,
-                                        k,
-                                        v,
-                                        attn_metadata,
-                                        attention_mask,
-                                        attention_window_size,
-                                        attention_mask_data,
-                                        mrope_config=mrope_config,
-                                        attention_sinks=attention_sinks,
-                                        has_lora=bool(lora_params))
+        attn_output = self.forward_impl(
+            q,
+            k,
+            v,
+            attn_metadata,
+            attention_mask,
+            attention_window_size,
+            attention_mask_data,
+            mrope_config=mrope_config,
+            attention_sinks=attention_sinks,
+            has_lora=bool(lora_params),
+            multi_item_part_lens=multi_item_part_lens,
+        )
 
         if self.attn_output_gate:
             gate = torch.sigmoid(gate)
@@ -1658,12 +1686,14 @@ def _get_attn_scale(position_ids: torch.Tensor) -> torch.Tensor:
         q = (q * attn_scale).to(q.dtype)
         return q
 
-    def forward_impl(self,
-                     position_ids: Optional[torch.Tensor],
-                     hidden_states: torch.Tensor,
-                     attn_metadata: AttentionMetadata,
-                     output: torch.Tensor,
-                     latent_cache_gen: Optional[torch.Tensor] = None) -> None:
+    def forward_impl(
+        self,
+        position_ids: Optional[torch.Tensor],
+        hidden_states: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+        output: torch.Tensor,
+        latent_cache_gen: Optional[torch.Tensor] = None,
+    ) -> None:
         """
         Forward pass for the MLA module. Writes result into output tensor in-place.
 

tensorrt_llm/_torch/modules/attention.py

@@ -43,7 +43,7 @@ def __init__(self, model_engine, dist):
 
         self.model_engine.warmup_encoder()
 
-    def batch_forward(self, inputs: Dict[str, Any]) -> Dict[str, torch.Tensor]:
+    def batch_forward(self, inputs: Dict[str, Any], **kwargs) -> Dict[str, torch.Tensor]:
         """Execute a pre-formed batch in one forward pass.
 
         Args:
@@ -54,7 +54,7 @@ def batch_forward(self, inputs: Dict[str, Any]) -> Dict[str, torch.Tensor]:
         Returns:
             Dict with 'logits' tensor and any other model outputs.
         """
-        return self.model_engine.encoder_forward(inputs)
+        return self.model_engine.encoder_forward(inputs, **kwargs)
 
     def shutdown(self):
         """No background thread to stop — just release model engine resources."""

tensorrt_llm/_torch/pyexecutor/encoder_executor.py

@@ -4514,7 +4514,8 @@ def _capture_encoder_cuda_graphs(self) -> None:
     @torch.inference_mode()
     @with_model_extra_attrs(lambda self: self.model.extra_attrs)
     @nvtx_range("encoder_forward")
-    def encoder_forward(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
+    def encoder_forward(self, inputs: Dict[str, Any],
+                        **kwargs) -> Dict[str, Any]:
         """Direct tensor-level forward for encode-only path.
 
         Bypasses ScheduledRequests/LlmRequest entirely. Takes a raw inputs
@@ -4547,32 +4548,40 @@ def encoder_forward(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
                 padded_inputs,
                 attn_metadata=graph_attn_metadata,
                 padded_num_tokens=key[1] if key is not None else None)
+            forward_kwargs = {
+                "gather_ids": None,
+                "gather_context_logits": False,
+                **kwargs,
+            }
 
             with with_shared_pool(
                     self.encoder_cuda_graph_runner.get_graph_pool()):
                 if key is None:
                     with MoeLoadBalancerIterContext(moe_load_balancer):
                         # Eager path — no graph for this bucket.
                         return self._forward_step(model_inputs,
-                                                  gather_ids=None,
-                                                  gather_context_logits=False)
+                                                  **forward_kwargs)
 
                 if self.encoder_cuda_graph_runner.needs_capture(key):
 
                     def forward_fn(
                             capture_inputs: Dict[str, Any]) -> Dict[str, Any]:
+                        capture_inputs = capture_inputs.copy()
+                        forward_kwargs = capture_inputs.pop("_forward_kwargs")
                         with MoeLoadBalancerIterContext(moe_load_balancer):
-                            return self._forward_step(
-                                capture_inputs,
-                                gather_ids=None,
-                                gather_context_logits=False)
+                            return self._forward_step(capture_inputs,
+                                                      **forward_kwargs)
 
                     self.encoder_cuda_graph_runner.capture(
-                        key, forward_fn, model_inputs)
+                        key, forward_fn, {
+                            **model_inputs, "_forward_kwargs": forward_kwargs
+                        })
 
                 with MoeLoadBalancerIterContext(moe_load_balancer):
                     graph_outputs = self.encoder_cuda_graph_runner.replay(
-                        key, model_inputs)
+                        key, {
+                            **model_inputs, "_forward_kwargs": forward_kwargs
+                        })
 
             # Return a clone to avoid sharing data_ptr with the static buffers.
             outputs = {}

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -66,6 +66,24 @@ class TokensPrompt(TypedDict):
     query_token_ids: NotRequired[List[int]]
     """The query input token IDs for star attention."""
 
+    multi_item_part_lens: NotRequired[List[int]]
+    """Metadata for multi-item scoring.
+
+    Specifying this enables the multi-item scoring mode.
+    In this mode, `prompt_token_ids` is assumed to consist of a prefix (or query) and
+    multiple items (or candidates). Each part of the prompt is assumed to be followed
+    by a delimiter token. The attention computation (masking and positional encodings)
+    is altered such that the prefill logits correspond to those that would
+    have been obtained upon running separate prefills for prefix + item1, prefix + item2,
+    and so on. Currently, this feature is only supported in prefill-only execution
+    via `LLM.encode`.
+
+    The list elements should correspond, in order, to prefix length and the length of
+    each following item to score (all lengths in tokens).
+    See https://github.com/flashinfer-ai/flashinfer/pull/1015 for details and nomenclature.
+    The length should not include the delimiter tokens.
+    """
+
 
 PromptInputs = Union[str, List[int], TextPrompt, TokensPrompt]