Use runtime token widths for Triton speculative verification (#39859)

Co-authored-by: raghotham <853234+raghotham@users.noreply.github.com>
This commit is contained in:
Lianmin Zheng
2026-09-18 11:07:55 -07:00
committed by GitHub
co-authored by raghotham
parent 6bd1a0af1d
commit 248c202b46
10 changed files with 302 additions and 65 deletions
@@ -16,6 +16,7 @@ import torch
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.srt.runtime_context import get_context
from sglang.srt.speculative.adaptive_runtime_state import SpecRuntimeState
from sglang.srt.speculative.eagle_info import EagleVerifyInput
from sglang.srt.speculative.eagle_utils import organize_draft_results
from sglang.srt.speculative.eagle_worker_v2 import EagleDraftWorker, EAGLEWorkerV2
from sglang.test.ci.ci_register import register_amd_ci, register_cpu_ci
@@ -135,6 +136,19 @@ class TestEagleWorkerV2Topk1FastPath(CustomTestCase):
with self.assertRaises(AssertionError):
worker._rebuild_topk1_chain_buffers()
def test_idle_verify_input_keeps_required_layout_tensors(self):
verify_input = EagleVerifyInput.create_idle_input(
topk=1,
spec_steps=3,
num_verify_tokens=4,
device=DEVICE,
)
self.assertEqual(verify_input.custom_mask.dtype, torch.bool)
self.assertEqual(verify_input.custom_mask.shape, (0,))
self.assertEqual(verify_input.positions.dtype, torch.int64)
self.assertEqual(verify_input.positions.shape, (0,))
def test_idle_draft_runs_each_eager_forward_without_tree_layout(self):
worker = object.__new__(EagleDraftWorker)
worker.speculative_num_steps = 3