Use runtime token widths for Triton speculative verification (#39859)
Co-authored-by: raghotham <853234+raghotham@users.noreply.github.com>
This commit is contained in:
co-authored by
raghotham
parent
6bd1a0af1d
commit
248c202b46
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user