wehub-resource-sync
76 lines
3.2 KiB
Markdown
76 lines
3.2 KiB
Markdown
# Stream SDK
|
|
|
|
> Stateful, single-request orchestration over the [KV Cache SDK](kvcache.md).
|
|
|
|
## Goal
|
|
|
|
`lmcache.sdk.kvcache` is stateless — it moves KV by token ids and forgets. A
|
|
token dropping request spans several passes (prefill, modify the cached KV,
|
|
decode), which are encoded as different requests. `LMCacheStream` binds all
|
|
request belonging to a request as a stream instead of leaving the management
|
|
to the user.
|
|
|
|
```py
|
|
import lmcache.sdk.kvcache as lmc_sdk
|
|
import lmcache.sdk.stream as lmc_stream
|
|
|
|
ctx = lmc_sdk.connect(url=..., http_url=..., model_name="Qwen/Qwen3-8B")
|
|
stream = lmc_stream.create_request(
|
|
ctx, post_completion, prompt_token_ids=source_tokens
|
|
)
|
|
|
|
stream.generate({"max_tokens": 1}) # prefill -> offload the prompt KV
|
|
stream.modify_kv(drop_tokens) # drop_tokens: retrieve -> edit -> store
|
|
stream.generate({"max_tokens": 256}) # replay the uncached tail + decode
|
|
```
|
|
|
|
`post_completion(prompt_token_ids, sampling_params, cache_salt)` calls the
|
|
engine (e.g. vLLM `/v1/completions`, streaming) and yields one
|
|
`TokenEvent(token_id, text)` per token. See the
|
|
[token-dropping example](../../../examples/token_dropping/).
|
|
|
|
Stream is not used on its own. `LMCacheBatchedStream` wraps it for submitting
|
|
batched requests. See [batch.md](batch.md).
|
|
|
|
## State model
|
|
|
|
The full logical sequence is `tokens` + `_suffix_tokens`. The token-bearing
|
|
fields mean different things:
|
|
|
|
- **`tokens`** — sequence backing the stored KV; sent as the next prompt.
|
|
Replaced when the KV is modified via `update_kv`.
|
|
- **`_suffix_tokens`** — tokens *not* in the stored KV: the non-chunk-aligned
|
|
tail left after `modify_kv`. Managed by the next `generate`.
|
|
- **`_decoded` / `_text_parts`** — cumulative generated count / text across all
|
|
`generate`s (`decoded_tokens()`, `output_text()`); unaffected by compaction.
|
|
- **`done`** — True once a `generate` yields `< max_tokens` (EOS). Reset when
|
|
`update_kv` is called.
|
|
|
|
## Suffix contract
|
|
|
|
`retrieve_kv` returns only a **chunk-aligned** prefix, so after `modify_kv`
|
|
edits it the remainder (sub-chunk tail + not-yet-offloaded tokens) has no
|
|
stored KV. The stream carries it across the edit:
|
|
|
|
- **`modify_kv`** records the uncached tail `tokens[cached_len:]` after KV
|
|
modification (e.g. compaction).
|
|
- **`generate`** prepends `_suffix_tokens` (plus any caller `suffix_tokens`)
|
|
to `tokens` for request submission, then clears it.
|
|
|
|
## Public API
|
|
|
|
- **`LMCacheStream(ctx, post_completion, prompt_token_ids, cache_salt="")`**
|
|
or the `create_request()`.
|
|
- **`generate(sampling_params, suffix_tokens=())`** returns `StreamPerfMetrics`
|
|
the performance metrics of one request
|
|
- **`modify_kv(fn)`**: calls `retrieve_kv`, passing the KV and tokens so far
|
|
to the function `fn(kv[2,L,T,D], tokens[:cached_len])` which returns the new
|
|
pair of `(new_kv, new_tokens)`, then replaces the original KV by `update_kv`.
|
|
- **`retrieve_kv(timeout=30, poll_interval=0.2)`** and
|
|
**`update_kv(kv, tokens)`**: wraps `lmc_sdk.retrieve()` and `lmc_sdk.store()`
|
|
- Accessors: `stream_id()`, `suffix_tokens()`, `decoded_tokens()`,
|
|
`output_text()`, `output_tokens()`, `is_done()`.
|
|
- **`StreamPerfMetrics`** (per call metric, times in **seconds**).
|
|
- **`TokenEvent(token_id, text)`**, **`PostCompletion`** (Protocol),
|
|
**`LMCacheStreamError`**.
|