[OPSD] Request last-position logits only in HybridEngine rollout prefill

deepspeed/runtime/rollout/hybrid_engine_rollout.py

@@ -780,6 +780,7 @@ def _continuous_prefill_cache(self, module, static_cache_type, prompt_ids, promp
             use_cache=True,
             cache_position=torch.arange(prompt_len, device=device),
             position_ids=position_ids,
+            **self._last_logits_kwargs(module),
         )
         return prefill_output.logits[:, -1, :].argmax(dim=-1, keepdim=True), prefill_cache
 
@@ -813,6 +814,15 @@ def _is_eos(self, token):
         eos_ids = torch.as_tensor(eos_token_id, dtype=token.dtype, device=token.device).flatten()
         return bool((token == eos_ids).any().item())
 
+    @staticmethod
+    def _last_logits_kwargs(module):
+        """Ask models that support it for last-position logits only, since prefill samples just one token."""
+        # Without this the prefill materializes [batch, prompt_len, vocab] logits, the largest activation of the
+        # rollout. Models that do not declare the argument may forward unknown kwargs to submodules, so skip them.
+        if "logits_to_keep" in signature(module.forward).parameters:
+            return {"logits_to_keep": 1}
+        return {}
+
     @staticmethod
     def _call_model(module, input_ids, **kwargs):
         """Call models with only the cache arguments their HF version accepts."""
@@ -990,6 +1000,7 @@ def _generate_graph(self, prompt_ids, prompt_attn, max_new_tokens, pad_token_id,
             past_key_values=prefill_cache,
             use_cache=True,
             cache_position=torch.arange(prompt_len, device=device),
+            **self._last_logits_kwargs(module),
         )
         next_token = prefill_out.logits[:, -1, :].argmax(dim=-1, keepdim=True)
 

tests/unit/runtime/rollout/test_hybrid_engine_rollout.py

@@ -265,6 +265,56 @@ def forward(self, input_ids, attention_mask, past_key_values=None, use_cache=Tru
     assert (1, 3) in model.calls
 
 
+def test_continuous_prefill_requests_last_logits_only():
+
+    class LastLogitsModel(torch.nn.Module):
+        _supports_cache_class = True
+
+        def __init__(self):
+            super().__init__()
+            self.weight = torch.nn.Parameter(torch.zeros(1))
+            self.prefill_logits_to_keep = []
+
+            self.config = PretrainedConfig(
+                max_position_embeddings=32,
+                num_hidden_layers=1,
+                num_attention_heads=1,
+                num_key_value_heads=1,
+                hidden_size=1,
+                head_dim=1,
+            )
+
+        # Mirrors the Transformers contract: 0 returns logits for every position, k for the last k.
+        def forward(self,
+                    input_ids,
+                    attention_mask,
+                    past_key_values=None,
+                    use_cache=True,
+                    cache_position=None,
+                    position_ids=None,
+                    logits_to_keep=0):
+            key_states = input_ids[:, None, :, None].to(dtype=torch.float32)
+            past_key_values.update(key_states,
+                                   key_states,
+                                   layer_idx=0,
+                                   cache_kwargs={"cache_position": cache_position})
+            if input_ids.shape[1] > 1:
+                self.prefill_logits_to_keep.append(logits_to_keep)
+            # Each position predicts its own token + 5, so reading any position but the last changes the output.
+            logits = torch.nn.functional.one_hot((input_ids + 5) % 16, num_classes=16).to(torch.float32)
+            if logits_to_keep:
+                logits = logits[:, -logits_to_keep:]
+            return SimpleNamespace(logits=logits, past_key_values=past_key_values)
+
+    model = LastLogitsModel()
+    rollout = HybridEngineRollout(SimpleNamespace(module=model), SimpleNamespace(pad_token_id=0, eos_token_id=2))
+    request = RolloutRequest(torch.tensor([[1, 2, 3], [1, 2, 4]]), torch.ones((2, 3), dtype=torch.long))
+    output = rollout.generate(request, SamplingConfig(max_new_tokens=2, temperature=0, continuous_batch_size=2))
+
+    assert output.input_ids[:, 3:].tolist() == [[8, 13], [9, 14]]
+    assert model.prefill_logits_to_keep == [1]
+
+
 def test_aligned_continuous_generation_supports_mixed_effective_prompt_lengths():
 
     class CacheConfig(SimpleNamespace):