chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Tiny extract ModelRunnerOutput (#15400)

Browse Source
This commit is contained in:
fzyzcjy
2025-12-18 22:18:45 +08:00
committed by GitHub
parent c5f4e20f2f
commit ad9616f13a
7 changed files with 34 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/speculative/eagle_worker.py
+5 -5
View File
@@ -646,9 +646,9 @@ class EAGLEWorker(TpModelWorker):
spec_info.hidden_states = hidden_states
# Run forward
logits_output, _ = self.draft_model_runner.forward(
logits_output = self.draft_model_runner.forward(
forward_batch, skip_attn_backend_init=True
)
).logits_output
if self.server_args.enable_nan_detection:
detect_nan(logits_output)
probs = torch.softmax(logits_output.next_token_logits, dim=-1)
@@ -946,7 +946,7 @@ class EAGLEWorker(TpModelWorker):
model_worker_batch, self.draft_model_runner
)
forward_batch.return_logprob = False
logits_output, _ = self.draft_model_runner.forward(forward_batch)
logits_output = self.draft_model_runner.forward(forward_batch).logits_output
if self.enable_nan_detection:
detect_nan(logits_output)
assert isinstance(forward_batch.spec_info, EagleDraftInput)
@@ -1023,9 +1023,9 @@ class EAGLEWorker(TpModelWorker):
self.draft_model_runner.attn_backend.init_forward_metadata(
forward_batch
)
logits_output, _ = self.draft_model_runner.forward(
logits_output = self.draft_model_runner.forward(
forward_batch, skip_attn_backend_init=True
)
).logits_output
self.capture_for_decode(logits_output, forward_batch.spec_info)
if self.enable_nan_detection:
python/sglang/srt/speculative/eagle_worker_v2.py
+5 -5
View File
@@ -391,9 +391,9 @@ class EagleDraftWorker(BaseDraftWorker):
spec_info.hidden_states = hidden_states
# Run forward
logits_output, _ = self.draft_runner.forward(
logits_output = self.draft_runner.forward(
forward_batch, skip_attn_backend_init=True
)
).logits_output
if self.server_args.enable_nan_detection:
detect_nan(logits_output)
probs = torch.softmax(logits_output.next_token_logits, dim=-1)
@@ -465,7 +465,7 @@ class EagleDraftWorker(BaseDraftWorker):
# Run forward
forward_batch = ForwardBatch.init_new(batch, self.draft_runner)
logits_output, _ = self.draft_runner.forward(forward_batch)
logits_output = self.draft_runner.forward(forward_batch).logits_output
# Update spec_info for the next draft step
probs = torch.softmax(logits_output.next_token_logits, dim=-1)
@@ -516,9 +516,9 @@ class EagleDraftWorker(BaseDraftWorker):
forward_batch
)
else:
draft_logits_output, _ = self.draft_runner.forward(
draft_logits_output = self.draft_runner.forward(
forward_batch, skip_attn_backend_init=True
)
).logits_output
# Reorganize the spec info for the next batch
draft_logits_output.next_token_logits = draft_logits_output.next_token_logits[