[DLLM] Add initial cuda graph support (#14203)

2025-12-08 14:12:35 +08:00
parent 661e9775d0
commit 36361adcbf
6 changed files with 93 additions and 8 deletions
--- a/python/sglang/srt/layers/logits_processor.py
+++ b/python/sglang/srt/layers/logits_processor.py
@@ -392,6 +392,14 @@ class LogitsProcessor(nn.Module):
                input_ids, hidden_states, lm_head, logits_metadata, multi_item_delimiter
            )

+        if logits_metadata.forward_mode.is_dllm_extend():
+            assert self.return_full_logits
+            full_logits = self._get_logits(hidden_states, lm_head, logits_metadata)
+            return LogitsProcessorOutput(
+                full_logits=full_logits,
+                next_token_logits=None,
+            )
+
        # Get the last hidden states and last logits for the next token prediction
        if (
            logits_metadata.forward_mode.is_decode_or_idle()