diff --git a/CMakeLists.txt b/CMakeLists.txt index 2592c0eef712..24916123051b 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -325,6 +325,17 @@ endif() # cumem_allocator extension # +if(VLLM_GPU_LANG STREQUAL "CUDA") + define_extension_target( + _ple_memops + DESTINATION vllm + LANGUAGE CXX + SOURCES "csrc/ple_memops.cpp" + LIBRARIES CUDA::cuda_driver + USE_SABI 3.8 + WITH_SOABI) +endif() + set(VLLM_CUMEM_EXT_SRC "csrc/cumem_allocator.cpp") diff --git a/csrc/ple_memops.cpp b/csrc/ple_memops.cpp new file mode 100644 index 000000000000..c9110e30fb3e --- /dev/null +++ b/csrc/ple_memops.cpp @@ -0,0 +1,67 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// Adapted flag protocol from FreeToken's ple_store_ext.cpp at +// af71ba43206e124f5ff6419b47ee36c6e9981078 (Apache-2.0). +// This implementation uses the CUDA driver directly and translates pinned +// host addresses rather than assuming identical host/device virtual addresses. +#include +#include +#include + +static CUresult device_flag(unsigned long long host, CUdeviceptr* device) { + return cuMemHostGetDevicePointer(device, reinterpret_cast(host), 0); +} + +static PyObject* signal_flag(PyObject*, PyObject* args) { + unsigned long long flag; + if (!PyArg_ParseTuple(args, "K", &flag)) return nullptr; + __atomic_store_n(reinterpret_cast(flag), uint64_t{1}, + __ATOMIC_RELEASE); + Py_RETURN_NONE; +} + +static PyObject* memop_write(PyObject*, PyObject* args) { + unsigned long long stream, host, value; + if (!PyArg_ParseTuple(args, "KKK", &stream, &host, &value)) return nullptr; + CUdeviceptr flag; + CUresult status = device_flag(host, &flag); + if (status == CUDA_SUCCESS) + status = cuStreamWriteValue64(reinterpret_cast(stream), flag, + value, CU_STREAM_WRITE_VALUE_DEFAULT); + return PyLong_FromLong(status); +} + +static PyObject* memop_wait_geq(PyObject*, PyObject* args) { + unsigned long long stream, host, value; + if (!PyArg_ParseTuple(args, "KKK", &stream, &host, &value)) return nullptr; + CUdeviceptr flag; + CUresult status = device_flag(host, &flag); + if (status == CUDA_SUCCESS) + status = cuStreamWaitValue64(reinterpret_cast(stream), flag, + value, CU_STREAM_WAIT_VALUE_GEQ); + return PyLong_FromLong(status); +} + +static PyObject* memop_wait_reset(PyObject*, PyObject* args) { + unsigned long long stream, host; + if (!PyArg_ParseTuple(args, "KK", &stream, &host)) return nullptr; + CUdeviceptr flag; + CUresult status = device_flag(host, &flag); + auto s = reinterpret_cast(stream); + if (status == CUDA_SUCCESS) + status = cuStreamWaitValue64(s, flag, 1, CU_STREAM_WAIT_VALUE_GEQ); + if (status == CUDA_SUCCESS) + status = cuStreamWriteValue64(s, flag, 0, CU_STREAM_WRITE_VALUE_DEFAULT); + return PyLong_FromLong(status); +} + +static PyMethodDef methods[] = { + {"signal_flag", signal_flag, METH_VARARGS, "Release-store the host flag."}, + {"memop_write", memop_write, METH_VARARGS, "Queue a 64-bit flag write."}, + {"memop_wait_geq", memop_wait_geq, METH_VARARGS, "Queue a flag wait."}, + {"memop_wait_reset", memop_wait_reset, METH_VARARGS, + "Queue wait and reset."}, + {nullptr, nullptr, 0, nullptr}}; +static PyModuleDef module = {PyModuleDef_HEAD_INIT, "_ple_memops", nullptr, -1, + methods}; +PyMODINIT_FUNC PyInit__ple_memops() { return PyModule_Create(&module); } diff --git a/setup.py b/setup.py index 81ff499be4d5..f829c6e8acba 100644 --- a/setup.py +++ b/setup.py @@ -1360,6 +1360,7 @@ def _read_requirements(filename: str) -> list[str]: ext_modules.append(CMakeExtension(name="vllm._rocm_C")) if _is_cuda(): + ext_modules.append(CMakeExtension(name="vllm._ple_memops")) ext_modules.append(CMakeExtension(name="vllm.vllm_flash_attn._vllm_fa2_C")) if USE_PRECOMPILED_EXTENSIONS or ( CUDA_HOME and get_nvcc_cuda_version() >= Version("12.3") diff --git a/tests/models/qwen4_exp/test_ple_deferred_cuda.py b/tests/models/qwen4_exp/test_ple_deferred_cuda.py new file mode 100644 index 000000000000..d6161cc8ef69 --- /dev/null +++ b/tests/models/qwen4_exp/test_ple_deferred_cuda.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Real memop/copy capture smoke; checkpoint integration is a separate gate.""" + +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import numpy as np +import pytest +import torch + +from vllm.config import CUDAGraphMode +from vllm.models.qwen4_exp.nvidia import ple_layer +from vllm.models.qwen4_exp.nvidia.ple_wait import ( + DeferredRows, + StreamMemopsUnavailable, +) + + +class _Table: + def __init__(self, source): + self.source = source + + def gather(self, ids): + # Production mmap gather consumes a flattened ID array. + assert ids.ndim == 1 + return np.stack( + [self.source[int(ids[h]), h].view(torch.uint8).numpy() for h in range(2)] + ) + + +def test_complete_passes_flat_ids_to_cuda_smoke_table(): + source = torch.arange(64, dtype=torch.bfloat16).reshape(8, 2, 4) + helper = object.__new__(DeferredRows) + helper.table = _Table(source) + helper.ids = torch.tensor([[5, 3]], dtype=torch.int64) + helper.rows = torch.empty((1, 2, 4), dtype=torch.bfloat16) + helper._poisoned = False + helper._pending = True + helper._readback_event = Mock() + helper._ext = Mock() + helper.flag = torch.zeros(1, dtype=torch.int64) + + helper.complete() + + expected = torch.stack([source[5, 0], source[3, 1]])[None] + assert torch.equal(helper.rows.view(torch.uint8), expected.view(torch.uint8)) + helper._readback_event.synchronize.assert_called_once_with() + helper._ext.signal_flag.assert_called_once_with(helper.flag.data_ptr()) + assert not helper.pending + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_runtime_none_capture_replays_fresh_rows_and_resets_flag(): + # Distinct raw BF16 rows catch stale data, zeros and ID mixups without + # relying on floating point comparisons or a model's generated text. + source = torch.arange(64, dtype=torch.bfloat16).reshape(8, 2, 4) + + destination = torch.empty((1, 2, 4), dtype=torch.bfloat16, device="cuda") + try: + helper = DeferredRows(destination, _Table(source)) + except StreamMemopsUnavailable as exc: + pytest.skip(str(exc)) + helper.prepare_dummy() + graph = torch.cuda.CUDAGraph() + context = SimpleNamespace( + cudagraph_runtime_mode=CUDAGraphMode.NONE, + no_compile_layers={ + "ple": SimpleNamespace(ple_embedding=SimpleNamespace(deferred_rows=helper)) + }, + ) + with ( + patch.object(ple_layer, "get_forward_context", return_value=context), + torch.cuda.graph(graph), + ): + ple_layer.qwen4_exp_ple_deferred_rows(destination, "ple") + + for row_ids in ([1, 2], [5, 3], [0, 7]): + ids = torch.tensor([row_ids], dtype=torch.int64, device="cuda") + helper.prepare(ids) + graph.replay() + # Never synchronize a replay waiting for the host before releasing it. + try: + helper.complete() + except BaseException: + helper.abort() + raise + torch.accelerator.synchronize() + expected = torch.stack([source[row_ids[h], h] for h in range(2)])[None] + assert torch.equal( + destination.cpu().view(torch.uint8), expected.view(torch.uint8) + ) + assert int(helper.flag[0]) == 0 + assert not helper.pending diff --git a/tests/models/qwen4_exp/test_ple_deferred_state.py b/tests/models/qwen4_exp/test_ple_deferred_state.py new file mode 100644 index 000000000000..19535810c208 --- /dev/null +++ b/tests/models/qwen4_exp/test_ple_deferred_state.py @@ -0,0 +1,161 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CPU lifecycle contracts for the model state.""" + +import unittest +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import torch + +from vllm.config import CUDAGraphMode +from vllm.models.qwen4_exp.nvidia import ple_layer +from vllm.models.qwen4_exp.nvidia.model_state import Qwen4ExpModelState +from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState + + +class DeferredStateTests(unittest.TestCase): + def test_consume_records_during_full_capture_runtime_none(self): + for mode, capturing, rows, expected in ( + (CUDAGraphMode.NONE, True, 1, True), + (CUDAGraphMode.NONE, False, 1, False), + (CUDAGraphMode.PIECEWISE, True, 1, False), + (CUDAGraphMode.FULL, False, 1, True), + (CUDAGraphMode.NONE, True, 2, False), + ): + with self.subTest(mode=mode, capturing=capturing, rows=rows): + helper = Mock() + context = SimpleNamespace( + cudagraph_runtime_mode=mode, + no_compile_layers={ + "ple": SimpleNamespace( + ple_embedding=SimpleNamespace(deferred_rows=helper) + ) + }, + ) + output = torch.empty(rows, 2, 3) + with ( + patch.object( + ple_layer, "get_forward_context", return_value=context + ), + patch.object( + torch.cuda, + "is_current_stream_capturing", + return_value=capturing, + ), + ): + ple_layer.qwen4_exp_ple_deferred_rows(output, "ple") + if expected: + helper.consume.assert_called_once_with(destination=output) + else: + helper.consume.assert_not_called() + + def make_state(self, count=3): + state = object.__new__(Qwen4ExpModelState) + state._mmap_ple_modules = tuple( + SimpleNamespace(deferred_rows=Mock()) for _ in range(count) + ) + state._deferred_ple_step = False + state._deferred_ple_poisoned = False + return state + + def test_mixed_capability_disables_all_helpers_before_capture(self): + state = self.make_state() + state.device = torch.device("cpu") + state.max_num_tokens = 8 + modules = state._mmap_ple_modules + modules[1].deferred_rows = None + for module in modules: + module.mmap_staging_nbytes = Mock(return_value=16) + module.initialize_mmap_staging = Mock() + with patch( + "vllm.models.qwen4_exp.nvidia.model_state.MemorySnapshot", + return_value=SimpleNamespace(free_memory=1024), + ): + state._initialize_mmap_staging(modules) + self.assertTrue(all(m.deferred_rows is None for m in modules)) + for module in modules: + module.initialize_mmap_staging.assert_called_once_with(8, state.device) + state.set_deferred_ple_step(True) + self.assertFalse(state._deferred_ple_step) + + def test_complete_all_layers_and_next_step(self): + state = self.make_state() + state.set_deferred_ple_step(True) + state.complete_deferred_ple() + for module in state._mmap_ple_modules: + module.deferred_rows.complete.assert_called_once() + module.deferred_rows.abort.assert_not_called() + self.assertFalse(state._deferred_ple_step) + state.set_deferred_ple_step(True) + self.assertTrue(state._deferred_ple_step) + + def test_fill_failure_releases_unvisited_layers_and_poison_latches(self): + state = self.make_state() + state._mmap_ple_modules[0].deferred_rows.complete.side_effect = ValueError( + "disk" + ) + state.set_deferred_ple_step(True) + with self.assertRaisesRegex(ValueError, "disk"): + state.complete_deferred_ple() + for module in state._mmap_ple_modules: + module.deferred_rows.abort.assert_called_once() + state._mmap_ple_modules[1].deferred_rows.complete.assert_not_called() + with self.assertRaisesRegex(RuntimeError, "poisoned"): + state.set_deferred_ple_step(False) + + def test_failed_release_still_attempts_every_layer(self): + state = self.make_state() + state.set_deferred_ple_step(True) + state._mmap_ple_modules[0].deferred_rows.abort.side_effect = ValueError( + "release" + ) + with self.assertRaisesRegex(ValueError, "release"): + state.abort_deferred_ple() + for module in state._mmap_ple_modules: + module.deferred_rows.abort.assert_called_once() + self.assertTrue(state._deferred_ple_poisoned) + + def test_prepare_failure_releases_already_prepared_layers(self): + state = self.make_state() + state.uses_ngram_embedding = True + state.ple_query_start_loc = torch.zeros(2, dtype=torch.int32) + state._prepare_ngram_context = Mock(return_value=torch.zeros((1, 2))) + batch = SimpleNamespace( + num_reqs_after_padding=1, + num_tokens=1, + num_tokens_after_padding=1, + num_reqs=1, + input_ids=torch.tensor([3]), + query_start_loc=torch.tensor([0, 1], dtype=torch.int32), + ) + for module in state._mmap_ple_modules: + module.prepare_deferred_mmap_rows = Mock() + state._mmap_ple_modules[1].prepare_deferred_mmap_rows.side_effect = ValueError( + "ids" + ) + state.set_deferred_ple_step(True) + with ( + patch.object(MambaHybridModelState, "prepare_inputs", return_value={}), + self.assertRaisesRegex(ValueError, "ids"), + ): + state.prepare_inputs(batch, None) + state._mmap_ple_modules[0].prepare_deferred_mmap_rows.assert_called_once() + state._mmap_ple_modules[2].prepare_deferred_mmap_rows.assert_not_called() + for module in state._mmap_ple_modules: + module.deferred_rows.abort.assert_called_once() + + def test_disabled_and_empty_have_no_effect(self): + for state in (self.make_state(), self.make_state(0)): + state.set_deferred_ple_step(False) + state.complete_deferred_ple() + state.abort_deferred_ple() + self.assertFalse(state._deferred_ple_poisoned) + state = self.make_state() + state._mmap_ple_modules[1].deferred_rows = None + state.set_deferred_ple_step(True) + self.assertFalse(state._deferred_ple_step) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/models/qwen4_exp/test_ple_wait.py b/tests/models/qwen4_exp/test_ple_wait.py new file mode 100644 index 000000000000..496ae6759ad2 --- /dev/null +++ b/tests/models/qwen4_exp/test_ple_wait.py @@ -0,0 +1,138 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CPU contracts for the deferred PLE host staging helper.""" + +from __future__ import annotations + +import unittest +from typing import Any +from unittest.mock import Mock, patch + +import torch + +from vllm.models.qwen4_exp.nvidia import ple_wait + + +class _FakeEvent: + def __init__(self, calls: list[str]) -> None: + self.calls = calls + + def synchronize(self) -> None: + self.calls.append("event") + + +class _FakeStream: + def __init__(self, calls: list[str]) -> None: + self.calls = calls + + def synchronize(self) -> None: + self.calls.append("stream") + + +class _FakeExtension: + def __init__(self, calls: list[str]) -> None: + self.calls = calls + + def signal_flag(self, _flag_ptr: int) -> None: + self.calls.append("signal") + + +def _bare_helper(rows: torch.Tensor) -> Any: + helper = object.__new__(ple_wait.DeferredRows) + helper.rows = rows + return helper + + +class DeferredRowsTests(unittest.TestCase): + def test_probe_unsupported_wait_fences_queued_write(self) -> None: + helper = _bare_helper(torch.empty(1)) + helper._ext = Mock() + helper._ext.memop_write.return_value = 0 + helper._ext.memop_wait_geq.return_value = 801 + stream = Mock(cuda_stream=12) + scratch = torch.zeros(1, dtype=torch.int64) + with ( + patch.object(ple_wait.torch, "zeros", return_value=scratch), + self.assertRaises(ple_wait.StreamMemopsUnavailable), + ): + helper._probe_stream_memops(stream) + stream.synchronize.assert_called_once_with() + + def test_probe_driver_error_is_not_capability_fallback(self) -> None: + helper = _bare_helper(torch.empty(1)) + helper._ext = Mock() + helper._ext.memop_write.return_value = 1 + stream = Mock(cuda_stream=12) + scratch = torch.zeros(1, dtype=torch.int64) + with ( + patch.object(ple_wait.torch, "zeros", return_value=scratch), + self.assertRaises(RuntimeError) as caught, + ): + helper._probe_stream_memops(stream) + self.assertNotIsInstance(caught.exception, ple_wait.StreamMemopsUnavailable) + helper._ext.memop_wait_geq.assert_not_called() + stream.synchronize.assert_called_once_with() + + def test_raw_uint8_rows_reinterpret_for_bfloat16_and_fp8(self) -> None: + for dtype in (torch.bfloat16, torch.float8_e4m3fn): + rows = torch.empty((1, 2, 3), dtype=dtype) + raw = torch.arange( + rows.numel() * rows.element_size(), dtype=torch.uint8 + ).reshape(2, -1) + actual = _bare_helper(rows)._as_rows_tensor(raw) + expected = raw.reshape(-1).view(dtype).reshape(rows.shape) + self.assertTrue( + torch.equal(actual.view(torch.uint8), expected.view(torch.uint8)) + ) + + def test_raw_uint8_rows_validate_byte_count(self) -> None: + rows = torch.empty((1, 2, 3), dtype=torch.bfloat16) + helper = _bare_helper(rows) + with self.assertRaisesRegex(ValueError, "bytes"): + helper._as_rows_tensor(torch.zeros((2, 5), dtype=torch.uint8)) + + def test_abort_fences_queued_reset_before_signal(self) -> None: + calls: list[str] = [] + helper = _bare_helper(torch.zeros((1, 1, 2), dtype=torch.bfloat16)) + helper._ext = _FakeExtension(calls) + helper.flag = torch.zeros(1, dtype=torch.int64) + helper._readback_event = _FakeEvent(calls) + helper._prepare_stream = _FakeStream(calls) + helper._pending = True + helper._rows_ready = False + helper._gate_armed = False + helper._reset_queued = True + helper._readback_recorded = True + helper._poisoned = False + helper._poison_reason = None + + helper.abort() + + self.assertEqual(calls, ["event", "signal"]) + self.assertTrue(helper.poisoned) + self.assertFalse(helper.pending) + self.assertFalse(helper._reset_queued) + + def test_abort_fences_prepare_partial_failure_before_signal(self) -> None: + calls: list[str] = [] + helper = _bare_helper(torch.zeros((1, 1, 2), dtype=torch.bfloat16)) + helper._ext = _FakeExtension(calls) + helper.flag = torch.zeros(1, dtype=torch.int64) + helper._readback_event = _FakeEvent(calls) + helper._prepare_stream = _FakeStream(calls) + helper._pending = False + helper._rows_ready = True + helper._gate_armed = False + helper._reset_queued = True + helper._readback_recorded = False + helper._poisoned = False + helper._poison_reason = None + + helper.abort() + + self.assertEqual(calls, ["stream", "signal"]) + self.assertTrue(helper.poisoned) + + +if __name__ == "__main__": + unittest.main() diff --git a/vllm/models/qwen4_exp/nvidia/model_state.py b/vllm/models/qwen4_exp/nvidia/model_state.py index ade850a7e524..41345ecb14d5 100644 --- a/vllm/models/qwen4_exp/nvidia/model_state.py +++ b/vllm/models/qwen4_exp/nvidia/model_state.py @@ -35,6 +35,8 @@ def __init__( config = self.model_config.hf_text_config self.uses_ngram_embedding = bool(config.ple_layer_ids) self._mmap_ple_modules: tuple[Qwen4ExpNGramEmbedding, ...] = () + self._deferred_ple_step = False + self._deferred_ple_poisoned = False if not self.uses_ngram_embedding: self.ngram_context_len = 0 self.ngram_eos_token_id = 0 @@ -157,6 +159,32 @@ def _initialize_mmap_staging( ) for module in modules: module.initialize_mmap_staging(self.max_num_tokens, self.device) + deferred = [m.deferred_rows for m in modules if m.deferred_rows is not None] + if deferred and len(deferred) != len(modules): + # Capture must agree with step eligibility across all mmap layers. + # Otherwise a capable layer would capture WAIT/H2D while the model + # uses synchronous preparation (which never signals its flag). + # Initialization precedes all prepare/capture calls, so no reader + # or pending producer owns these helpers yet. + for module in modules: + module.deferred_rows = None + logger.warning( + "PLE deferred unavailable on some mmap layers; " + "using synchronous preparation on every mmap layer" + ) + deferred = [] + if deferred: + pinned_bytes = sum( + tensor.numel() * tensor.element_size() + for rows in deferred + for tensor in (rows.ids, rows.rows, rows.flag) + ) + logger.info( + "PLE mmap deferred enabled: %d layers, pinned_bytes=%d; " + "FULL graph with one real token only, eager/prefill unchanged", + len(deferred), + pinned_bytes, + ) def _dummy_query_start_loc_and_context( self, num_reqs: int, num_tokens: int @@ -215,6 +243,46 @@ def _prepare_ngram_context( ) return context + def set_deferred_ple_step(self, eligible: bool) -> None: + if self._deferred_ple_poisoned: + raise RuntimeError("Deferred PLE is poisoned after a failed fill") + self._deferred_ple_step = bool( + eligible + and self._mmap_ple_modules + and all(m.deferred_rows is not None for m in self._mmap_ple_modules) + ) + + def abort_deferred_ple(self) -> None: + if not self._deferred_ple_step: + return + # Release every layer: a failed early fill must not strand a later + # captured WAIT. The model is permanently unusable after this point. + self._deferred_ple_poisoned = True + self._deferred_ple_step = False + first_error: BaseException | None = None + for module in self._mmap_ple_modules: + if module.deferred_rows is not None: + try: + module.deferred_rows.abort() + except BaseException as error: + if first_error is None: + first_error = error + if first_error is not None: + raise first_error + + def complete_deferred_ple(self) -> None: + if not self._deferred_ple_step: + return + try: + for module in self._mmap_ple_modules: + assert module.deferred_rows is not None + module.deferred_rows.complete() + except BaseException: + self.abort_deferred_ple() + raise + finally: + self._deferred_ple_step = False + def prepare_inputs( self, input_batch: InputBatch, @@ -245,14 +313,26 @@ def prepare_inputs( actual_input_ids = input_batch.input_ids[:actual_tokens] actual_query_start_loc = query_start_loc[: num_reqs + 1] actual_ngram_context = ngram_context[:num_reqs] - for module in self._mmap_ple_modules: - module.prepare_mmap_rows( - actual_input_ids, - actual_query_start_loc, - actual_ngram_context, - actual_tokens, - padded_tokens, - ) + try: + for module in self._mmap_ple_modules: + if self._deferred_ple_step: + module.prepare_deferred_mmap_rows( + actual_input_ids, + actual_query_start_loc, + actual_ngram_context, + ) + else: + module.prepare_mmap_rows( + actual_input_ids, + actual_query_start_loc, + actual_ngram_context, + actual_tokens, + padded_tokens, + ) + except BaseException: + if self._deferred_ple_step: + self.abort_deferred_ple() + raise return model_inputs def prepare_dummy_inputs( diff --git a/vllm/models/qwen4_exp/nvidia/ple_layer.py b/vllm/models/qwen4_exp/nvidia/ple_layer.py index 04d20288e2f9..59c5349fcc3d 100644 --- a/vllm/models/qwen4_exp/nvidia/ple_layer.py +++ b/vllm/models/qwen4_exp/nvidia/ple_layer.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """GPU-resident Qwen4Exp position-learning enhancement layers.""" +import os from collections.abc import Iterable, Sequence from typing import cast @@ -11,6 +12,7 @@ from vllm.compilation.breakable_cudagraph import eager_break_during_capture from vllm.config import CacheConfig, ModelConfig, VllmConfig, get_current_vllm_config +from vllm.config.compilation import CUDAGraphMode from vllm.forward_context import get_forward_context from vllm.model_executor.layers.linear import MergedColumnParallelLinear from vllm.model_executor.layers.mamba.abstract import MambaBase @@ -40,7 +42,7 @@ from vllm.transformers_utils.configs.qwen4_exp import ( Qwen4ExpTextConfig, ) -from vllm.utils.torch_utils import get_dtype_size +from vllm.utils.torch_utils import direct_register_custom_op, get_dtype_size from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum from vllm.v1.attention.backends.short_conv_attn import ( PleShortConvAttentionBackend, @@ -50,6 +52,7 @@ from ..common.ple import PLEVocabParallelEmbedding from . import ple_mmap from .ops.ple import ple_conv, ple_gate, ple_ngram_ids +from .ple_wait import DeferredRows, StreamMemopsUnavailable class Qwen4ExpPLEGroupedNorm(nn.Module): @@ -327,6 +330,7 @@ def __init__( # staging (see initialize_mmap_staging). None until V2 model state # allocates it; stays None for the non-mmap embedding. self._mmap_staging: torch.Tensor | None = None + self.deferred_rows: DeferredRows | None = None if ple_mmap.enabled(): vllm_config = get_current_vllm_config() ple_mmap.check_cudagraph_safety(vllm_config) @@ -522,6 +526,32 @@ def initialize_mmap_staging( device=device, ) + if os.getenv("VLLM_PLE_MMAP_DEFERRED", "0") == "1": + table = self._require_mmap_embedding().table + if table is None: + raise RuntimeError("Deferred PLE requires an initialized mmap table") + try: + self.deferred_rows = DeferredRows(self._mmap_staging[:1], table) + except StreamMemopsUnavailable as exc: + import warnings + + warnings.warn( + f"Deferred PLE unavailable; using synchronous mmap rows: {exc}", + RuntimeWarning, + stacklevel=2, + ) + + def prepare_deferred_mmap_rows( + self, + input_ids: torch.Tensor, + query_start_loc: torch.Tensor, + ngram_context: torch.Tensor, + ) -> None: + if self.deferred_rows is None: + raise RuntimeError("Deferred PLE was not initialized") + ids = self.compute_ngram_ids(input_ids, query_start_loc, ngram_context) + self.deferred_rows.prepare(ids) + def prepare_mmap_rows( self, input_ids: torch.Tensor, @@ -563,6 +593,8 @@ def prepare_dummy_mmap_rows(self, padded_tokens: int) -> None: f"PLE mmap: {self.layer_name!r} staging was never initialized" ) self._mmap_staging[:padded_tokens].zero_() + if padded_tokens == 1 and self.deferred_rows is not None: + self.deferred_rows.prepare_dummy() def forward( self, @@ -582,7 +614,10 @@ def forward( # query_start_loc.size()[0] under the old whole-forward custom # op. A plain shape[0] read stays a SymInt when traced. num_tokens = input_ids.reshape(-1).shape[0] - return self._mmap_staging[:num_tokens].flatten(-2) + output = self._mmap_staging[:num_tokens] + if self.deferred_rows is not None: + torch.ops.vllm.qwen4_exp_ple_deferred_rows(output, self.layer_name) + return output.flatten(-2) if query_start_loc is None or ngram_context is None: raise RuntimeError("PLE inputs were not prepared") ngram_ids = self.compute_ngram_ids(input_ids, query_start_loc, ngram_context) @@ -1045,3 +1080,31 @@ def forward( "Qwen4ExpPLEGroupedNorm", "Qwen4ExpPLELayer", ] + + +def qwen4_exp_ple_deferred_rows(output: torch.Tensor, layer_name: str) -> None: + """Capture stream WAIT/H2D; host completes fills after graph dispatch.""" + context = get_forward_context() + capture = ( + context.cudagraph_runtime_mode == CUDAGraphMode.NONE + and torch.cuda.is_current_stream_capturing() + ) + # FULL capture invokes the model with runtime mode NONE. PIECEWISE and + # ordinary eager execution must keep the synchronous staging path. + if ( + context.cudagraph_runtime_mode == CUDAGraphMode.FULL or capture + ) and output.shape[0] == 1: + layer = context.no_compile_layers[layer_name] + layer.ple_embedding.deferred_rows.consume(destination=output) + + +def qwen4_exp_ple_deferred_rows_fake(output: torch.Tensor, layer_name: str) -> None: + return + + +direct_register_custom_op( + op_name="qwen4_exp_ple_deferred_rows", + op_func=qwen4_exp_ple_deferred_rows, + mutates_args=["output"], + fake_impl=qwen4_exp_ple_deferred_rows_fake, +) diff --git a/vllm/models/qwen4_exp/nvidia/ple_wait.py b/vllm/models/qwen4_exp/nvidia/ple_wait.py new file mode 100644 index 000000000000..968e2608ef62 --- /dev/null +++ b/vllm/models/qwen4_exp/nvidia/ple_wait.py @@ -0,0 +1,517 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Deferred host staging for the Qwen4Exp PLE row lookup. + +The captured side of a PLE lookup is a fixed-shape copy from pinned host +memory. This module keeps the request-dependent row IDs and gathered rows in +stable pinned buffers, and uses CUDA stream memops to let a captured +consumer wait for the host gather without a device synchronization. + +The flag protocol and extension API are reused from FreeToken's Apache-2.0 +implementation at +``python/freetoken/kernel/csrc/ple_store/ple_store_ext.cpp`` (FreeToken +commit ``af71ba43206e124f5ff6419b47ee36c6e9981078``). The extension is +loaded lazily so importing vLLM on a CPU-only host does not require CUDA or a +separate FreeToken build. +""" + +from __future__ import annotations + +import importlib +from contextlib import suppress +from typing import Any + +import torch + +_EXTENSION_NAME = "vllm._ple_memops" +_EXTENSION_CALLS = ( + "memop_wait_reset", + "memop_write", + "memop_wait_geq", + "signal_flag", +) +_extension: Any | None = None + + +class StreamMemopsUnavailable(RuntimeError): + """The optional CUDA stream-memop capability is unavailable at startup.""" + + +def _load_extension() -> Any: + """Load the optional stream-memop extension with an actionable error.""" + global _extension + if _extension is not None: + return _extension + try: + module = importlib.import_module(_EXTENSION_NAME) + except (ImportError, OSError) as exc: + raise StreamMemopsUnavailable( + f"Qwen4Exp deferred PLE requires a CUDA vLLM build with {_EXTENSION_NAME}." + ) from exc + missing = [ + name for name in _EXTENSION_CALLS if not callable(getattr(module, name, None)) + ] + if missing: + raise RuntimeError( + f"{_EXTENSION_NAME} is missing required callables: {', '.join(missing)}" + ) + _extension = module + return module + + +def _stream_for(device: torch.device, stream: Any | None) -> Any: + """Return a CUDA stream, accepting the torch stream wrapper used by vLLM.""" + if stream is not None: + return stream + return torch.cuda.current_stream(device) + + +def _stream_ptr(stream: Any) -> int: + """Get the CUDA stream handle accepted by the in-tree CUDA extension.""" + try: + return int(stream.cuda_stream) + except AttributeError: + return int(stream) + + +def _check_memop_status(status: Any, operation: str) -> None: + """Report a rejected CUDA driver operation without hiding its status.""" + if status is not None and int(status) != 0: + raise RuntimeError( + f"PLE stream operation {operation} failed with CUDA status {status}" + ) + + +class DeferredRows: + """Stage one fixed-shape PLE row batch across a captured forward. + + ``destination`` is the stable device staging tensor consumed by the + captured model. It must be shaped ``[1, heads, head_dim]`` and remain + alive for the lifetime of this object. ``table`` must provide + ``gather(np.ndarray)`` and return one row per flattened ID. The table is + intentionally kept on the host side: no table read or allocation occurs + during capture. + + The normal sequence is:: + + rows.prepare(ids_cuda) + dispatch_forward() + rows.complete() + rows.consume() # called by the captured forward + + ``prepare_dummy`` supplies zero rows and a pre-signaled flag for warmup or + graph capture. ``abort`` unblocks a pending consumer before poisoning the + object, so an exception cannot leave a CUDA graph waiting forever. + """ + + def __init__( + self, + destination: torch.Tensor, + table: Any, + *, + stream: Any | None = None, + ) -> None: + if not isinstance(destination, torch.Tensor): + raise TypeError("DeferredRows destination must be a torch.Tensor") + if destination.device.type != "cuda": + raise ValueError("DeferredRows destination must be a CUDA tensor") + if destination.ndim != 3 or destination.shape[0] != 1: + raise ValueError( + "DeferredRows destination must have shape [1, heads, head_dim]" + ) + if any(int(size) <= 0 for size in destination.shape[1:]): + raise ValueError("DeferredRows destination dimensions must be positive") + if not destination.is_contiguous(): + raise ValueError("DeferredRows destination must be contiguous") + if not callable(getattr(table, "gather", None)): + raise TypeError("DeferredRows table must provide gather(ids_numpy)") + + # Resolve this before allocating the pinned buffers. An absent opt-in + # extension should fail at construction, rather than much later from + # inside a graph replay. + self._ext = _load_extension() + self.destination = destination + self.table = table + self._stream_override = stream + self.stream = _stream_for(destination.device, stream) + shape = tuple(int(size) for size in destination.shape) + self.ids = torch.empty( + (1, shape[1]), dtype=torch.int64, device="cpu", pin_memory=True + ) + self.rows = torch.empty( + shape, dtype=destination.dtype, device="cpu", pin_memory=True + ) + self.flag = torch.empty((1,), dtype=torch.int64, device="cpu", pin_memory=True) + self.ids.zero_() + self.rows.zero_() + # A freshly allocated helper represents a dummy-ready staging buffer. + # This is required for a warmup/capture path that consumes before a + # real host gather has completed. + self.flag.fill_(1) + self._probe_stream_memops(self.stream) + self._readback_event = torch.cuda.Event() + self._pending = False + self._rows_ready = True + self._gate_armed = False + self._reset_queued = False + self._readback_recorded = False + self._prepare_stream: Any | None = None + self._poisoned = False + self._poison_reason: BaseException | None = None + + @classmethod + def allocate( + cls, + destination: torch.Tensor, + table: Any, + *, + stream: Any | None = None, + ) -> DeferredRows: + """Construct a helper, retaining the stable destination address.""" + return cls(destination, table, stream=stream) + + @property + def pending(self) -> bool: + """Whether a prepared ID batch still needs :meth:`complete`.""" + return self._pending + + @property + def poisoned(self) -> bool: + """Whether a failed operation permanently disabled this helper.""" + return self._poisoned + + @property + def poison_reason(self) -> BaseException | None: + """Return the first exception that poisoned this helper, if any.""" + return self._poison_reason + + @property + def flag_ptr(self) -> int: + """Pinned flag address for diagnostics and integration tests.""" + return int(self.flag.data_ptr()) + + @property + def ids_pinned(self) -> torch.Tensor: + """Pinned host ID readback buffer.""" + return self.ids + + @property + def rows_pinned(self) -> torch.Tensor: + """Pinned host row buffer.""" + return self.rows + + def _ensure_healthy(self) -> None: + if self._poisoned: + detail = f": {self._poison_reason}" if self._poison_reason else "" + raise RuntimeError(f"DeferredRows is permanently poisoned{detail}") + + def _validate_ids(self, ids: torch.Tensor) -> None: + if not isinstance(ids, torch.Tensor): + raise TypeError("DeferredRows IDs must be a torch.Tensor") + if ids.device != self.destination.device: + raise ValueError( + f"DeferredRows IDs device {ids.device} != " + f"destination device {self.destination.device}" + ) + expected = (1, int(self.destination.shape[1])) + if tuple(ids.shape) != expected: + raise ValueError( + f"DeferredRows IDs shape {tuple(ids.shape)} != expected {expected}" + ) + if ids.dtype != torch.int64: + raise ValueError(f"DeferredRows IDs must be torch.int64, got {ids.dtype}") + if not ids.is_contiguous(): + raise ValueError("DeferredRows IDs must be contiguous") + + def _validate_destination(self, destination: torch.Tensor) -> None: + if not isinstance(destination, torch.Tensor): + raise TypeError("DeferredRows destination must be a torch.Tensor") + if destination is not self.destination: + if destination.device != self.destination.device: + raise ValueError( + f"DeferredRows destination device {destination.device} != " + f"{self.destination.device}" + ) + if tuple(destination.shape) != tuple(self.destination.shape): + raise ValueError( + f"DeferredRows destination shape {tuple(destination.shape)} " + f"!= {tuple(self.destination.shape)}" + ) + if destination.dtype != self.destination.dtype: + raise ValueError( + f"DeferredRows destination dtype {destination.dtype} != " + f"{self.destination.dtype}" + ) + if not destination.is_contiguous(): + raise ValueError("DeferredRows destination must be contiguous") + + def _signal(self) -> None: + """Signal the host flag, accepting only the extension's API.""" + self._ext.signal_flag(self.flag_ptr) + + def _operation_stream(self, stream: Any | None = None) -> Any: + """Resolve an operation stream, honoring an optional constructor pin.""" + selected = self._stream_override if stream is None else stream + resolved = _stream_for(self.destination.device, selected) + self.stream = resolved + return resolved + + def _probe_stream_memops(self, stream: Any) -> None: + """Probe outside capture, retaining scratch until queued writes finish.""" + scratch = torch.zeros((1,), dtype=torch.int64, device="cpu", pin_memory=True) + stream_ptr = _stream_ptr(stream) + try: + for operation, function in ( + ("write", self._ext.memop_write), + ("wait", self._ext.memop_wait_geq), + ): + status = function(stream_ptr, int(scratch.data_ptr()), 7) + # CUDA_ERROR_NOT_SUPPORTED. Other failures are not a safe + # reason to silently select a different execution path. + if status == 801: + raise StreamMemopsUnavailable( + f"CUDA stream memop {operation} is not supported" + ) + _check_memop_status(status, f"{operation}(probe, 7)") + finally: + stream.synchronize() + if int(scratch[0]) != 7: + raise RuntimeError( + "FreeToken PLE stream memop probe did not publish its value" + ) + + def prepare(self, ids: torch.Tensor) -> None: + """Queue the fixed-shape ID readback before dispatching the forward. + + The stream write of zero is deliberately queued before the D2H copy. + A host ``flag.zero_`` would race an earlier graph replay and can leave + the next replay waiting on a value that belongs to the wrong batch. + """ + self._ensure_healthy() + if self._pending: + raise RuntimeError("DeferredRows.prepare called while a batch is pending") + self._validate_ids(ids) + stream = self._operation_stream() + try: + status = self._ext.memop_write(_stream_ptr(stream), self.flag_ptr, 0) + _check_memop_status(status, "memop_write(flag, 0)") + self._prepare_stream = stream + self._reset_queued = True + self.ids.copy_(ids, non_blocking=True) + self._readback_event.record(stream) + self._readback_recorded = True + self._pending = True + self._rows_ready = False + self._gate_armed = False + except BaseException as exc: + self._poison(exc) + raise + + def _as_rows_tensor(self, gathered: Any) -> torch.Tensor: + """Convert a table result to the destination's dtype and fixed shape.""" + if isinstance(gathered, torch.Tensor): + source = gathered.detach() + if source.device.type != "cpu": + raise ValueError( + "DeferredRows table must return host rows; CUDA table " + "results are not copied during completion" + ) + else: + # MmapPleTable.gather returns a writable NumPy uint8 array. The + # conversion intentionally stays on the host and avoids a second + # device allocation during the deferred completion callback. + source = torch.as_tensor(gathered, device="cpu") + source = source.contiguous() + target_nbytes = self.rows.numel() * self.rows.element_size() + source_nbytes = source.numel() * source.element_size() + if source_nbytes != target_nbytes: + raise ValueError( + f"DeferredRows table returned {source_nbytes} bytes, expected " + f"{target_nbytes}" + ) + # Mmap tables are raw bytes. Reinterpret them for every destination + # dtype, including BF16 where one row occupies twice as many bytes as + # its element count. A tensor table result is converted only when it + # already has the destination element count. + if source.dtype == torch.uint8: + source = source.reshape(-1).view(self.rows.dtype) + elif source.numel() == self.rows.numel(): + source = source.to(dtype=self.rows.dtype) + else: + raise ValueError( + "DeferredRows table returned a non-byte result with an " + "incompatible element count" + ) + return source.reshape(self.rows.shape) + + def complete(self) -> None: + """Finish the host gather and release the captured consumer.""" + self._ensure_healthy() + if not self._pending: + raise RuntimeError("DeferredRows.complete called without prepare") + try: + self._readback_event.synchronize() + ids = self.ids.numpy().reshape(-1) + gathered = self.table.gather(ids) + self.rows.copy_(self._as_rows_tensor(gathered)) + self._signal() + self._reset_queued = False + self._readback_recorded = False + self._prepare_stream = None + self._pending = False + self._rows_ready = True + self._gate_armed = False + except BaseException as exc: + self._poison(exc) + raise + + def gate( + self, + *, + stream: Any | None = None, + capture: bool | None = None, + dummy: bool = False, + ) -> None: + """Apply the flag operation for integrations that split wait/copy. + + A real captured path emits ``WAIT(>=1); RESET``. Eager and dummy + paths use a host signal because they do not need a device wait. The + regular :meth:`consume` method invokes this operation itself. + """ + self._ensure_healthy() + if self._pending: + raise RuntimeError("DeferredRows.gate called while a batch is pending") + if dummy: + self.prepare_dummy() + return + stream = self._operation_stream(stream) + if capture is None: + capture = bool(torch.cuda.is_current_stream_capturing()) + try: + if capture: + status = self._ext.memop_wait_reset(_stream_ptr(stream), self.flag_ptr) + _check_memop_status(status, "memop_wait_reset") + self._gate_armed = True + else: + self._signal() + self._gate_armed = False + except BaseException as exc: + self._poison(exc) + raise + + def consume( + self, + destination: torch.Tensor | None = None, + *, + stream: Any | None = None, + capture: bool | None = None, + wait: bool = True, + ) -> torch.Tensor: + """Copy the completed or dummy rows into the fixed device staging. + + Under CUDA graph capture, the wait/reset is emitted before the H2D + copy. ``wait=False`` is available when a caller already invoked + :meth:`gate` for a split implementation. + """ + self._ensure_healthy() + if self._pending: + raise RuntimeError("DeferredRows.consume called while a batch is pending") + if not self._rows_ready: + raise RuntimeError("DeferredRows.consume has no completed rows") + destination = self.destination if destination is None else destination + self._validate_destination(destination) + stream = self._operation_stream(stream) + if capture is None: + capture = bool(torch.cuda.is_current_stream_capturing()) + try: + if wait: + if self._gate_armed: + self._gate_armed = False + elif capture: + status = self._ext.memop_wait_reset( + _stream_ptr(stream), self.flag_ptr + ) + _check_memop_status(status, "memop_wait_reset") + else: + # Keep the eager path harmless after a caller reused a + # helper without a preceding captured wait. + self._signal() + destination.copy_(self.rows, non_blocking=True) + return destination + except BaseException as exc: + self._poison(exc) + raise + + def prepare_dummy(self) -> None: + """Reset pinned rows for warmup/capture after the prior copy finishes.""" + self._ensure_healthy() + if self._pending: + raise RuntimeError("DeferredRows.prepare_dummy called while pending") + try: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("DeferredRows.prepare_dummy cannot run in capture") + self._operation_stream().synchronize() + self.rows.zero_() + self._signal() + self._rows_ready = True + self._gate_armed = False + except BaseException as exc: + self._poison(exc) + raise + + def _poison(self, exc: BaseException) -> None: + """Unblock a possible wait, then permanently fail closed.""" + if self._poisoned: + return + # ``prepare`` queues the flag reset before recording the readback + # event. Do not signal the host flag ahead of that reset: a later + # queued reset could overwrite the release and strand a graph wait. + # The readback event is recorded before graph launch, so synchronizing + # it cannot wait on the captured consumer. If recording failed, + # prepare is still outside capture and synchronizing its stream is the + # only safe recovery fence. + if self._reset_queued: + try: + if self._readback_recorded: + self._readback_event.synchronize() + elif self._prepare_stream is not None: + self._prepare_stream.synchronize() + except BaseException: + # Preserve the original failure and still attempt the host + # release below. The object is poisoned either way. + pass + try: + self._signal() + except BaseException: + # The original exception is more useful to the caller. A best + # effort host store keeps a graph from waiting if the extension's + # signal wrapper itself failed after a partial operation. + with suppress(BaseException): + self.flag.fill_(1) + self._reset_queued = False + self._readback_recorded = False + self._prepare_stream = None + self._pending = False + self._rows_ready = False + self._gate_armed = False + self._poisoned = True + self._poison_reason = exc + + def abort(self) -> None: + """Signal any pending consumer and permanently poison this helper.""" + if self._poisoned: + return + self._poison(RuntimeError("DeferredRows aborted")) + + +# Keep the earlier coordination name available while the integration uses the +# more descriptive DeferredRows constructor. +PLEWait = DeferredRows +PLEDeferredStaging = DeferredRows + + +__all__ = [ + "DeferredRows", + "PLEDeferredStaging", + "PLEWait", +] diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 093d3c8b35c9..8f277f2108ca 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1762,6 +1762,16 @@ def execute_model( if inputs_embeds is not None and not requires_raw_input_tokens(self.model): input_ids = None + deferred_ple_setter = getattr(self.model_state, "set_deferred_ple_step", None) + if deferred_ple_setter is not None: + deferred_ple_setter( + not dummy_run + and batch_desc.cg_mode == CUDAGraphMode.FULL + and input_batch.num_tokens == 1 + and input_batch.num_tokens_after_padding == 1 + and input_batch.num_reqs == 1 + ) + model_inputs = { "input_ids": input_ids, "positions": input_batch.positions, @@ -1780,43 +1790,60 @@ def execute_model( else self.model_state.prepare_inputs(input_batch, self.req_states) ), } - if not self.is_first_pp_rank: - # Update for non-first PP ranks. - model_inputs["input_ids"] = None - model_inputs["inputs_embeds"] = None - - # Prepare the intermediate tensors. - assert intermediate_tensors is not None - assert self.intermediate_tensors is not None - n = input_batch.num_tokens_after_padding - new_tensors = { - k: v[:n] - if dummy_run - else v[:n].copy_(intermediate_tensors.tensors[k][:n]) - for k, v in self.intermediate_tensors.tensors.items() - } - model_inputs["intermediate_tensors"] = IntermediateTensors(new_tensors) - del intermediate_tensors - - # Update the EPLB meta. - ubatch_slices = ubatch_state.slices if ubatch_state is not None else None - self.eplb.prepare_forward( - self.model_config, input_batch.num_tokens, ubatch_slices - ) + try: + if not self.is_first_pp_rank: + # Update for non-first PP ranks. + model_inputs["input_ids"] = None + model_inputs["inputs_embeds"] = None + + # Prepare the intermediate tensors. + assert intermediate_tensors is not None + assert self.intermediate_tensors is not None + n = input_batch.num_tokens_after_padding + new_tensors = { + k: v[:n] + if dummy_run + else v[:n].copy_(intermediate_tensors.tensors[k][:n]) + for k, v in self.intermediate_tensors.tensors.items() + } + model_inputs["intermediate_tensors"] = IntermediateTensors(new_tensors) + del intermediate_tensors + + # Update the EPLB meta. + ubatch_slices = ubatch_state.slices if ubatch_state is not None else None + self.eplb.prepare_forward( + self.model_config, input_batch.num_tokens, ubatch_slices + ) - self.step_timing.record_batch( - input_batch, batch_desc.cg_mode == CUDAGraphMode.FULL - ) - self.step_timing.forward_start() + self.step_timing.record_batch( + input_batch, batch_desc.cg_mode == CUDAGraphMode.FULL + ) + self.step_timing.forward_start() + except BaseException: + abort_ple = getattr(self.model_state, "abort_deferred_ple", None) + if abort_ple is not None: + abort_ple() + raise # Run model. if batch_desc.cg_mode == CUDAGraphMode.FULL: # Use explicit cudagraph replay for FULL mode. # NOTE(woosuk): Here, we don't need to pass the input tensors, # because they are already copied to the CUDA graph input buffers. - assert self.cudagraph_manager is not None - self.kv_connector.pre_forward(scheduler_output) - model_output = self.cudagraph_manager.run_fullgraph(batch_desc) + try: + assert self.cudagraph_manager is not None + self.kv_connector.pre_forward(scheduler_output) + model_output = self.cudagraph_manager.run_fullgraph(batch_desc) + # Host fills must release graph WAITs before any observer, + # connector or output path can synchronize the compute stream. + complete_ple = getattr(self.model_state, "complete_deferred_ple", None) + if complete_ple is not None: + complete_ple() + except BaseException: + abort_ple = getattr(self.model_state, "abort_deferred_ple", None) + if abort_ple is not None: + abort_ple() + raise else: # For piecewise and eager mode, just call model(). batch_descriptor = BatchDescriptor(