Ai00-X / Ai00-X/ai00_server

Feature request: allow stateless inference flows

Aberta
#170 0 comentários 0 reações 0 responsáveis Ver no GitHub
Linguagem predominante
Rust
Estrelas
619
Forks
74
Métricas de merge de PRs
Nenhum PR com merge em 30d

Descrição

## 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%.

Guia de contribuição

Nenhum guia de contribuição indexado para este repositório

Direção de pesquisa

Start by reading the CompletionRequest and ChatRequest entry points for /oai/v1/chat/completions and /oai/v1/completions, then trace execution through run.rs and ai00-core. Review chat.rs for stateful message handling and the cache paths described in the issue. Done means request and response state support, no_cache behavior, final-state extraction for stop and length, and final-chunk SSE support are defined and covered across these flows.

Escrita pelo modelo de indexação a partir do texto da issue.

Avaliação

Stack de tecnologia
rust
Domínio
ai, api, backend
Tipo de issue
Funcionalidade
Dificuldade
5/5
Tempo estimado
Mais de uma semana
Status de atividade
Ativa
Clareza
Razoavelmente clara
Facilidade para iniciantes
38/100

Receba novas issues na sua caixa de entrada

Um resumo curto de issues do GitHub para quem está começando.