Ai00-X / Ai00-X/ai00_server

Feature request: allow stateless inference flows

Open
#170 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Rust
Stars
619
Forks
74
PR merge metrics
No merged PRs in 30d

Description

## What

Add request parameters to `/oai/v1/chat/completions` and `/oai/v1/completions` to allow fully stateless request-response flows:
1. `return_state: bool` (or `return_state: "no" | "base64" | "json"`) to return the final hidden state tensor in the response.
2. `no_cache: bool` (default: `false`) to bypass saving prefix states into the internal Trie and `CacheHub`.

## Why

RWKV's fixed-size $O(1)$ recurrent hidden state provides a unique superpower over Transformers ($O(N)$ KV-cache): hidden states can be exchanged over standard networks with fixed bandwidth.

Enabling stateless inference unlocks:
- **Horizontal scaling & load balancing:** Requests can be routed round-robin across GPU worker pools without sticky sessions or distributed Redis state caches.
- **Zero server RAM leaks:** Long-running inactive sessions consume zero host memory or GPU slot memory.
- **Client-managed branching & persistence:** Clients can store conversation checkpoints locally and branch conversations deterministically without server-side timeouts or eviction issues.

Currently, achieving this requires a round-trip to `/oai/states` combined with prefix-cache lookups, which requires multiple requests and still relies on the server-side cache.

## Proposed Implementation Details

### 1. Request Parameters (`CompletionRequest` and `ChatRequest`)
- `return_state: bool` (default: `false`)
- `no_cache: bool` (default: `false`)
- `state_encoding: Option` (`"base64"` (default) or `"json"`)

### 2. Response Format
Include an optional `state` field in `ChatResponse` and `CompletionChoice` (or root `CompletionResponse`):

```json
"state": {
"id": "00000000-0000-0000-0000-000000000000",
"shape": [1, 32, 4096, 1],
"data": "..." // Base64 byte string (or [f32] array if state_encoding="json")
}
```

*Note: The response shape should match what `InputState::Value` accepts so clients can pass it directly back into subsequent requests without transformation.*

### 3. Server-Side Execution (`run.rs` & `ai00-core`)
- **Trie Caching:** When `no_cache: true`, skip prompt cache registration in `process()` and completion cache insertion at `FinishReason::Stop`.
- **Memory Management (`caches.backed`):** When `no_cache: true`, ensure `InputState::Value` is loaded directly into the execution slot without permanently persisting inside `caches.backed: HashMap`, preventing host memory leaks.
- **State Extraction:** When `return_state: true`, call `self.back(batch)` upon reaching **both** `FinishReason::Stop` and `FinishReason::Length`, and pass the tensor through `context.sender`.

### 4. Handling Chat Context (`chat.rs`)
In stateless mode, clients will typically pass the previous `state` along with only the *new* turn message(s). We should clarify documentation or handle message templating so that historical turns are not re-tokenized and fed on top of an already-evaluated state.

### 5. SSE Streaming
For `stream: true`, embed the `state` field inside the final SSE chunk (accompanying `finish_reason`), maintaining compatibility with OpenAI-compatible SSE client parsers before sending `[DONE]`.

## Performance Considerations

- **Payload format:** An RWKV-6 3B state contains ~5.2M floats (~21 MB binary f32). In raw JSON, this results in ~55–65 MB of ASCII text, incurring heavy CPU serialization overhead. Using Base64 binary encoding reduces this to ~28 MB and drastically speeds up serialization/deserialization.
- **HTTP Compression:** Enabling gzip or zstd over HTTP should further compress the base64-encoded or binary state payload by at least 25%.

Contributor guide

No contributing guide indexed for this repository

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.