diff --git a/python/minisgl/core.py b/python/minisgl/core.py index be4d643e..0a70be81 100644 --- a/python/minisgl/core.py +++ b/python/minisgl/core.py @@ -40,6 +40,7 @@ def __post_init__(self) -> None: self.device_len = len(self.input_ids) self.max_device_len = len(self.input_ids) + self.output_len assert 0 <= self.cached_len < self.device_len <= self.max_device_len + self._host_buf: torch.Tensor | None = None @property def remain_len(self) -> int: @@ -54,7 +55,16 @@ def complete_one(self) -> None: self.device_len += 1 def append_host(self, next_token: torch.Tensor) -> None: - self.input_ids = torch.cat([self.input_ids, next_token]) + # Write into a lazily pre-allocated buffer (O(1) per token) instead of + # using torch.cat (O(n^2) per request) to prevent GPU stalling on the scheduler thread. + # Lazy allocation ensures non-decoding requests pay zero memory overhead. + # Previously written positions are read-only to guarantee radix cache safety. + if self._host_buf is None: + self._host_buf = torch.empty(self.max_device_len, dtype=self.input_ids.dtype) + self._host_buf[: len(self.input_ids)] = self.input_ids + new_len = len(self.input_ids) + 1 + self._host_buf[new_len - 1] = next_token + self.input_ids = self._host_buf[:new_len] @property def can_decode(self) -> bool: diff --git a/tests/core/test_req_append.py b/tests/core/test_req_append.py new file mode 100644 index 00000000..ff5334f3 --- /dev/null +++ b/tests/core/test_req_append.py @@ -0,0 +1,82 @@ +from __future__ import annotations + +import pytest +import torch +from minisgl.core import Req, SamplingParams + + +def _make_req(prompt_len: int = 8, output_len: int = 16, dtype=torch.int32) -> Req: + return Req( + input_ids=torch.arange(prompt_len, dtype=dtype), + table_idx=0, + cached_len=0, + output_len=output_len, + uid=0, + sampling_params=SamplingParams(max_tokens=output_len), + cache_handle=None, # type: ignore[arg-type] + ) + + +def test_append_host_matches_cat_reference(): + req = _make_req(prompt_len=8, output_len=16) + reference = req.input_ids.clone() + + for token in range(100, 116): + next_token = torch.tensor([token], dtype=torch.int32) + reference = torch.cat([reference, next_token]) + req.complete_one() + req.append_host(next_token) + assert req.input_ids.is_cpu + assert torch.equal(req.input_ids, reference) + assert len(req.input_ids) == req.device_len + + +def test_append_host_preserves_dtype(): + for dtype in (torch.int32, torch.int64): + req = _make_req(dtype=dtype) + req.complete_one() + req.append_host(torch.tensor([42], dtype=torch.int32)) + assert req.input_ids.dtype == dtype + + +def test_append_host_never_mutates_written_positions(): + # The radix cache stores slices of `input_ids` by reference (see + # RadixTreeNode.set_key_value), so previously-appended positions must + # never change under later appends. + req = _make_req(prompt_len=4, output_len=8) + req.complete_one() + req.append_host(torch.tensor([100], dtype=torch.int32)) + snapshot_view = req.input_ids[:5] # what a radix key would hold + snapshot_copy = snapshot_view.clone() + + for token in range(101, 108): + req.complete_one() + req.append_host(torch.tensor([token], dtype=torch.int32)) + + assert torch.equal(snapshot_view, snapshot_copy) + + +def test_append_host_overflow_raises(): + # Appending beyond max_device_len must fail loudly, not silently drop. + req = _make_req(prompt_len=2, output_len=1) + req.complete_one() + req.append_host(torch.tensor([7], dtype=torch.int32)) + req.complete_one() + with pytest.raises(IndexError): + req.append_host(torch.tensor([8], dtype=torch.int32)) + + +def test_chunked_req_append_still_forbidden(): + from minisgl.scheduler.prefill import ChunkedReq + + req = ChunkedReq( + input_ids=torch.arange(4, dtype=torch.int32), + table_idx=0, + cached_len=0, + output_len=4, + uid=0, + sampling_params=SamplingParams(), + cache_handle=None, # type: ignore[arg-type] + ) + with pytest.raises(NotImplementedError): + req.append_host(torch.tensor([1], dtype=torch.int32))