class DeepseekV41ModelState(DefaultModelState):
"""DefaultModelState plus the engram lookback window.
The engram n-gram hash needs the ids of the ``depth`` tokens preceding
each request's chunk start (see ``common/engram.py``). The runner keeps
the full token history on device, so the window is gathered there every
step: exact for prompt and generated tokens alike, whatever instance
produced their KV.
"""
def __init__(
self,
vllm_config: VllmConfig,
model: nn.Module,
encoder_cache: EncoderCache | None,
device: torch.device,
):
super().__init__(vllm_config, model, encoder_cache, device)
depth = model.token_lookback_depth
self.lookback_token_ids: torch.Tensor | None = None
if depth > 0:
# Persistent so a captured graph can read it on replay.
self.lookback_token_ids = torch.full(
(self.max_num_reqs, depth), -1, dtype=torch.int32, device=device
)
def prepare_inputs(
self, input_batch: InputBatch, req_states: RequestState
) -> dict[str, torch.Tensor | None]:
model_inputs = super().prepare_inputs(input_batch, req_states)
window = self.lookback_token_ids
if window is None:
return model_inputs
all_token_ids = req_states.all_token_ids.gpu
depth = window.shape[1]
_gather_lookback_kernel[(window.shape[0],)](
window,
input_batch.idx_mapping,
req_states.num_computed_tokens.gpu,
all_token_ids,
all_token_ids.stride(0),
input_batch.idx_mapping.shape[0],
DEPTH=depth,
BLOCK_DEPTH=triton.next_power_of_2(depth),
)
model_inputs["lookback_token_ids"] = window
return model_inputs
def prepare_dummy_inputs(self, num_reqs: int, num_tokens: int) -> dict[str, Any]:
model_inputs = super().prepare_dummy_inputs(num_reqs, num_tokens)
if self.lookback_token_ids is not None:
# The captured graph reads this buffer; replays refill it in place.
self.lookback_token_ids.fill_(-1)
model_inputs["lookback_token_ids"] = self.lookback_token_ids
return model_inputs