[Fix] Preserve model runner contracts in prefill CUDA graphs (#35452)
Co-authored-by: Oasis-Git <ayw.sirius19@gmail.com> Co-authored-by: Ke Bao <ispobaoke@gmail.com>
This commit is contained in:
Aurick Qiao
co-authored by
Oasis-Git
Ke Bao
Oasis-Git
Ke Bao
parent
745de73ba3
commit
d229952e25
@@ -1527,6 +1527,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
|
||||
column_starts=0,
|
||||
req_lens=shape_inputs["extend_seq_lens"],
|
||||
)
|
||||
forward_batch = self.model_runner.prepare_dummy_forward_batch(forward_batch)
|
||||
self.tbo_plugin.capture_one_batch_size(forward_batch, num_tokens=num_tokens)
|
||||
return forward_batch, self.model_runner.attn_backend
|
||||
|
||||
@@ -1860,6 +1861,9 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
|
||||
# The n-gram hasher runs outside the graph and reads this at replay.
|
||||
static_forward_batch.ngram_embedding_info = forward_batch.ngram_embedding_info
|
||||
static_forward_batch.engram_history = forward_batch.engram_history
|
||||
static_forward_batch = self.model_runner.prepare_dummy_forward_batch(
|
||||
static_forward_batch
|
||||
)
|
||||
if self._is_full_backend:
|
||||
forward_batch.next_token_logits_buffer = (
|
||||
static_forward_batch.next_token_logits_buffer
|
||||
@@ -2070,6 +2074,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
|
||||
input_token_ids_logprobs_val=output.input_token_ids_logprobs_val,
|
||||
input_token_ids_logprobs_idx=output.input_token_ids_logprobs_idx,
|
||||
input_logprobs_copy_done=output.input_logprobs_copy_done,
|
||||
customized_info=output.customized_info,
|
||||
mm_input_embeds=mm_input_embeds,
|
||||
)
|
||||
|
||||
|
||||
@@ -8,6 +8,7 @@ import torch
|
||||
|
||||
import sglang.srt.model_executor.model_runner_components.cuda_graph_setup as graph_setup
|
||||
import sglang.srt.model_executor.runner.prefill_cuda_graph_runner as runner_module
|
||||
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
|
||||
from sglang.srt.model_executor.cuda_graph_config import Backend
|
||||
from sglang.srt.model_executor.forward_batch_info import (
|
||||
CaptureHiddenMode,
|
||||
@@ -140,6 +141,79 @@ class TestPrefillCudaGraphRunnerChunkedPrefix(CustomTestCase):
|
||||
self.assertIs(runner.prepare_dummy_forward_batch(batch), batch)
|
||||
self.assertEqual(batch.attn_tp_sequence_sharded, expected)
|
||||
|
||||
def test_capture_prepare_applies_model_runner_batch_hook(self):
|
||||
class Slot:
|
||||
def __init__(self, buffer):
|
||||
self.buffer = buffer
|
||||
|
||||
def slice_for(self, _bs, num_tokens):
|
||||
return self.buffer[:num_tokens]
|
||||
|
||||
slots = {
|
||||
"input_ids": Slot(torch.zeros(4, dtype=torch.int64)),
|
||||
"out_cache_loc": Slot(torch.zeros(4, dtype=torch.int64)),
|
||||
"positions": Slot(torch.zeros(4, dtype=torch.int64)),
|
||||
}
|
||||
registry = SimpleNamespace(
|
||||
has_slot=lambda name: name in slots,
|
||||
get_slot=lambda name: slots[name],
|
||||
)
|
||||
prepared = []
|
||||
captured = []
|
||||
attention_backend = object()
|
||||
model_runner = SimpleNamespace(
|
||||
model_config=SimpleNamespace(context_len=8),
|
||||
pp_group=SimpleNamespace(is_last_rank=False),
|
||||
ngram_embedding_manager=SimpleNamespace(enabled=False),
|
||||
prepare_dummy_forward_batch=lambda batch: prepared.append(batch) or batch,
|
||||
attn_tp_sequence_sharded=lambda _: False,
|
||||
attn_backend=attention_backend,
|
||||
)
|
||||
runner = PrefillCudaGraphRunner.__new__(PrefillCudaGraphRunner)
|
||||
runner.model_runner = model_runner
|
||||
runner.device = torch.device("cpu")
|
||||
runner.prefill_backend_name = Backend.BREAKABLE
|
||||
runner.max_context_size = None
|
||||
runner._capture_req_slots = 1
|
||||
runner.max_bs = 4
|
||||
runner._prefill_static_buffers = None
|
||||
runner.buffer_registry = registry
|
||||
runner.require_mlp_tp_gather = False
|
||||
runner.require_attn_tp_gather = False
|
||||
runner._capture_lora = False
|
||||
runner.capture_hidden_mode = CaptureHiddenMode.NULL
|
||||
runner.static_draft_hidden_states = None
|
||||
runner.capture_return_pooled_hidden_states = False
|
||||
runner._next_token_logits_buffer = lambda _rows: None
|
||||
runner._build_capture_spec_info = lambda _num_tokens: None
|
||||
runner._capture_num_token_non_padded = lambda _num_tokens: None
|
||||
runner.tbo_plugin = SimpleNamespace(
|
||||
capture_one_batch_size=lambda batch, **_: captured.append(batch)
|
||||
)
|
||||
|
||||
batch, backend = runner.capture_prepare(4)
|
||||
|
||||
self.assertEqual(prepared, [batch])
|
||||
self.assertEqual(captured, [batch])
|
||||
self.assertIs(backend, attention_backend)
|
||||
|
||||
def test_trim_logits_output_preserves_customized_info(self):
|
||||
runner = PrefillCudaGraphRunner.__new__(PrefillCudaGraphRunner)
|
||||
runner.model_runner = SimpleNamespace(
|
||||
spec_algorithm=SimpleNamespace(is_speculative=lambda: False),
|
||||
)
|
||||
runner.raw_bs = 1
|
||||
runner._is_full_backend = True
|
||||
customized_info = {"per_request": [object()]}
|
||||
output = runner._trim_logits_output(
|
||||
LogitsProcessorOutput(
|
||||
next_token_logits=torch.zeros((4, 8)),
|
||||
customized_info=customized_info,
|
||||
)
|
||||
)
|
||||
|
||||
self.assertIs(output.customized_info, customized_info)
|
||||
|
||||
def test_low_free_memory_still_captures_prefill_graph(self):
|
||||
eager_runner = object()
|
||||
prefill_runner = object()
|
||||
@@ -247,7 +321,10 @@ class TestPrefillCudaGraphRunnerChunkedPrefix(CustomTestCase):
|
||||
runner.max_context_size = None
|
||||
runner._capture_chunked_prefix = False
|
||||
runner.buffer_registry = _FakeBatchRegistry()
|
||||
runner.model_runner = SimpleNamespace(attn_tp_sequence_sharded=lambda _: False)
|
||||
runner.model_runner = SimpleNamespace(
|
||||
attn_tp_sequence_sharded=lambda _: False,
|
||||
prepare_dummy_forward_batch=lambda batch: batch,
|
||||
)
|
||||
runner.enable_cp_bcg_capture = False
|
||||
runner._is_full_backend = False
|
||||
runner.backend = SimpleNamespace()
|
||||
|
||||
Reference in New Issue
Block a user