From dd42f0ec0fab9648a03640e69a471327dab9b902 Mon Sep 17 00:00:00 2001 From: "Jin, Youzhi" Date: Fri, 16 Oct 2026 15:23:35 +0000 Subject: [PATCH 1/2] Request last-position logits only in HybridEngine rollout prefill Prefill samples one token, but the static-cache and graph-capture paths called the model without logits_to_keep, so the lm_head produced [batch, prompt_len, vocab] logits. Pass logits_to_keep=1 when the model's forward declares it. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi --- .../runtime/rollout/hybrid_engine_rollout.py | 11 ++++ .../rollout/test_hybrid_engine_rollout.py | 50 +++++++++++++++++++ 2 files changed, 61 insertions(+) diff --git a/deepspeed/runtime/rollout/hybrid_engine_rollout.py b/deepspeed/runtime/rollout/hybrid_engine_rollout.py index f249cb84cd32..b71ca7295908 100644 --- a/deepspeed/runtime/rollout/hybrid_engine_rollout.py +++ b/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) diff --git a/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py b/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py index 84a96e645f37..de3bb64ca6a9 100644 --- a/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py +++ b/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) + logits = torch.zeros((input_ids.shape[0], input_ids.shape[1], 16)) + logits[:, :, 7] = 1 + 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() == [[7, 7], [7, 7]] + assert model.prefill_logits_to_keep == [1] + + def test_aligned_continuous_generation_supports_mixed_effective_prompt_lengths(): class CacheConfig(SimpleNamespace): From fba8f91844d1500a01f0a88380b0189339f19d64 Mon Sep 17 00:00:00 2001 From: "Jin, Youzhi" Date: Sat, 17 Oct 2026 15:24:02 +0000 Subject: [PATCH 2/2] Use position-dependent logits in the last-logits prefill test Every position used to predict the same token, so the test would also pass if the wrong position were sampled. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi --- tests/unit/runtime/rollout/test_hybrid_engine_rollout.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py b/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py index de3bb64ca6a9..961ac260ec79 100644 --- a/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py +++ b/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py @@ -300,8 +300,8 @@ def forward(self, cache_kwargs={"cache_position": cache_position}) if input_ids.shape[1] > 1: self.prefill_logits_to_keep.append(logits_to_keep) - logits = torch.zeros((input_ids.shape[0], input_ids.shape[1], 16)) - logits[:, :, 7] = 1 + # 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) @@ -311,7 +311,7 @@ def forward(self, 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() == [[7, 7], [7, 7]] + assert output.input_ids[:, 3:].tolist() == [[8, 13], [9, 14]] assert model.prefill_logits_to_keep == [1]