Xiaoze Fan
45 lines
1.4 KiB
Python
45 lines
1.4 KiB
Python
"""Req.append_host writes into the preallocated buffer: value-equivalent to the
|
|
old per-step torch.cat, no reallocation, and existing views stay stable."""
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from freetoken.core import Req, SamplingParams
|
|
from freetoken.scheduler.prefill import ChunkedReq
|
|
|
|
|
|
def _mk(cls, input_ids, output_len=4):
|
|
return cls(
|
|
input_ids=input_ids,
|
|
table_idx=0,
|
|
cached_len=0,
|
|
output_len=output_len,
|
|
uid=0,
|
|
sampling_params=SamplingParams(),
|
|
cache_handle=None,
|
|
)
|
|
|
|
|
|
def test_append_host_matches_cat_without_reallocating():
|
|
req = _mk(Req, torch.arange(6, dtype=torch.int32))
|
|
ref = torch.arange(6, dtype=torch.int32)
|
|
base_ptr = req.input_ids.data_ptr()
|
|
view = req.input_ids[:3]
|
|
for t in (101, 102, 103, 104):
|
|
tok = torch.tensor([t], dtype=torch.int32)
|
|
ref = torch.cat([ref, tok])
|
|
req.append_host(tok)
|
|
assert torch.equal(req.input_ids, ref)
|
|
assert req.input_ids.data_ptr() == base_ptr
|
|
assert req.input_ids.dtype == torch.int32
|
|
assert len(req.input_ids) == req.max_device_len
|
|
assert torch.equal(view, torch.arange(3, dtype=torch.int32))
|
|
|
|
|
|
def test_chunked_req_keeps_prompt_view_and_rejects_append():
|
|
ids = torch.arange(6, dtype=torch.int32)
|
|
req = _mk(ChunkedReq, ids)
|
|
assert req.input_ids is ids
|
|
with pytest.raises(NotImplementedError):
|
|
req.append_host(torch.tensor([1], dtype=torch.int32))
|