feat: add configurations

README.md

@@ -51,26 +51,29 @@ make build
 
 ```bash
 # Download a model
-puma pull inftyai/tiny-random-gpt2
+puma pull qwen/qwen2.5-0.5b
+
+# Run a model in an interactive chat
+puma run qwen/qwen2.5-0.5b
 
 # List all models
 puma ls
 
 # Inspect model details
-puma inspect inftyai/tiny-random-gpt2
+puma inspect qwen/qwen2.5-0.5b
 
 # Check system info
 puma info
 
 # Remove a model
-puma rm inftyai/tiny-random-gpt2
+puma rm qwen/qwen2.5-0.5b
 ```
 
 ### API Server
 
 ```bash
 # Start the inference server with a model
-puma serve inftyai/tiny-random-gpt2
+puma serve qwen/qwen2.5-0.5b
 
 # Server will start on http://0.0.0.0:8000
 # API endpoints:
@@ -91,7 +94,7 @@ curl http://localhost:8000/health
 curl http://localhost:8000/v1/chat/completions \
   -H "Content-Type: application/json" \
   -d '{
-    "model": "inftyai/tiny-random-gpt2",
+    "model": "qwen/qwen2.5-0.5b",
     "messages": [{"role": "user", "content": "Hello!"}]
   }'
 
@@ -109,9 +112,9 @@ curl http://localhost:8000/v1/chat/completions \
 | `rm <model>` | ✅ | Remove model and cache |
 | `info` | ✅ | Display system information |
 | `version` | ✅ | Show PUMA version |
+| `run <model>` | ✅ | Run a model in an interactive chat |
 | `serve <model>` | ✅ | Start OpenAI-compatible API server with a model |
 | `ps` | 🚧 | List running models |
-| `run` | 🚧 | Start model inference |
 | `stop` | 🚧 | Stop running model |
 
 ## Advanced Usage
@@ -144,6 +147,17 @@ puma ls llama -l author=meta
 
 **Available filters:** `author`, `task`, `license`, `provider`, `model_series`
 
+### Engine Tuning
+
+Both `run` and `serve` accept flags to tune the inference engine (KV-cache pool,
+block size, batch size, default token budget):
+
+```bash
+puma run qwen/qwen2.5-0.5b --max-batch-size 64 --default-max-tokens 256
+```
+
+See [docs/configuration.md](docs/configuration.md) for the full list of flags and defaults.
+
 ## API Server
 
 PUMA provides an OpenAI-compatible API server for model inference.
@@ -152,23 +166,25 @@ PUMA provides an OpenAI-compatible API server for model inference.
 
 ```bash
 # Start server with a model (default: 0.0.0.0:8000)
-puma serve inftyai/tiny-random-gpt2
+puma serve qwen/qwen2.5-0.5b
 
 # Custom host and port
-puma serve inftyai/tiny-random-gpt2 --host 127.0.0.1 --port 3000
+puma serve qwen/qwen2.5-0.5b --host 127.0.0.1 --port 3000
 
 # Model must be pulled first
-puma pull inftyai/tiny-random-gpt2
+puma pull qwen/qwen2.5-0.5b
 ```
 
+Engine parameters (KV-cache, batch size, etc.) can be tuned with additional flags — see [docs/configuration.md](docs/configuration.md).
+
 ### API Endpoints
 
 #### Chat Completions (Recommended)
 ```bash
 curl http://localhost:8000/v1/chat/completions \
   -H "Content-Type: application/json" \
   -d '{
-    "model": "inftyai/tiny-random-gpt2",
+    "model": "qwen/qwen2.5-0.5b",
     "messages": [
       {"role": "system", "content": "You are a helpful assistant."},
       {"role": "user", "content": "Hello!"}
@@ -183,7 +199,7 @@ curl http://localhost:8000/v1/chat/completions \
 curl http://localhost:8000/v1/chat/completions \
   -H "Content-Type: application/json" \
   -d '{
-    "model": "inftyai/tiny-random-gpt2",
+    "model": "qwen/qwen2.5-0.5b",
     "messages": [{"role": "user", "content": "Tell me a story"}],
     "stream": true
   }'
@@ -214,7 +230,7 @@ client = OpenAI(
 )
 
 response = client.chat.completions.create(
-    model="inftyai/tiny-random-gpt2",
+    model="qwen/qwen2.5-0.5b",
     messages=[
         {"role": "user", "content": "Hello!"}
     ]
@@ -226,28 +242,28 @@ print(response.choices[0].message.content)
 ### Inspect Output
 
 ```bash
-$ puma inspect inftyai/tiny-random-gpt2
+$ puma inspect qwen/qwen2.5-0.5b
 
-name: inftyai/tiny-random-gpt2
+name: qwen/qwen2.5-0.5b
 kind: Model
 spec:
-  author:         inftyai
-  model_series:   gpt2
+  author:         qwen
+  model_series:   qwen2
   task:           text-generation
-  license:        MIT
-  context_window: 2.05K
+  license:        APACHE-2.0
+  context_window: 32.77K
   safetensors:
-    total:        7.00B
+    total:        494.03M
     parameters:
-      f32:        7.00B
-  provider:     huggingface
+      bf16:        494.03M
+  provider:       huggingface
   cache:
-    revision:     abc123de
-    size:         1.24 GB
-    path:         ~/.puma/cache/...
+    revision:       060db6499f32faf8b98477b0a26969ef7d8b9987
+    size:           988.10 MB
+    path:           ~/.puma/cache/huggingface/models--qwen--qwen2.5-0.5b
 status:
-  created:      2 hours ago
-  updated:      2 hours ago
+  created:        2 months ago
+  updated:        2 months ago
 ```
 
 ## Model Management

docs/configuration.md

@@ -0,0 +1,36 @@
+# Configuration
+
+PUMA's inference engine exposes a set of tunable parameters. Both `puma run`
+and `puma serve` accept the same engine flags — all are optional and fall back
+to the defaults below.
+
+## Engine flags
+
+| Flag | Default | Description |
+|------|---------|-------------|
+| `--memory-pool-bytes` | `104857600` (100 MB) | Total KV-cache memory pool, in bytes |
+| `--block-size-bytes` | `512` | Size of a single KV block, in bytes |
+| `--tokens-per-block` | `16` | Number of tokens whose KV state fits in one block |
+| `--max-batch-size` | `32` | Maximum number of sequences batched together per step |
+| `--default-max-tokens` | `100` | Completion-token budget when a request omits `max_tokens` |
+
+## Examples
+
+```bash
+# Larger KV-cache pool and batch size, bigger default generation budget
+puma serve qwen/qwen2.5-0.5b \
+  --memory-pool-bytes 209715200 \
+  --max-batch-size 64 \
+  --default-max-tokens 256
+
+# The same flags work for interactive run
+puma run qwen/qwen2.5-0.5b --tokens-per-block 32
+```
+
+## Notes
+
+- **`--tokens-per-block`** must be greater than `0`.
+- **`--default-max-tokens`** only applies when a request omits `max_tokens`; an
+  explicit per-request `max_tokens` always takes precedence.
+- The flag defaults are sourced from `EngineConfig` in
+  `src/backend/llm_engine.rs`, so the CLI and the library stay in sync.

src/api/chat.rs

@@ -97,7 +97,8 @@ async fn chat_completions_non_stream(
         .generate(
             &req.model,
             &prompt,
-            req.max_tokens.unwrap_or(100),
+            req.max_tokens
+                .unwrap_or_else(|| engine.default_max_tokens()),
             req.temperature.unwrap_or(0.7),
         )
         .await?;
@@ -169,7 +170,8 @@ async fn chat_completions_stream(
             .generate_stream(
                 &model,
                 &prompt,
-                req.max_tokens.unwrap_or(100),
+                req.max_tokens
+                    .unwrap_or_else(|| engine.default_max_tokens()),
                 req.temperature.unwrap_or(0.7),
             )
             .await

src/api/completions.rs

@@ -74,7 +74,8 @@ pub async fn completions(
         .generate(
             &req.model,
             &prompt,
-            req.max_tokens.unwrap_or(100),
+            req.max_tokens
+                .unwrap_or_else(|| engine.default_max_tokens()),
             req.temperature.unwrap_or(0.7),
         )
         .await

src/api/tests.rs

@@ -13,8 +13,8 @@ use tempfile::TempDir;
 use tower::util::ServiceExt; // for `oneshot` and `ready`
 
 use super::routes::create_router;
-use crate::backend::engine;
 use crate::backend::mock::MockEngine;
+use crate::backend::{engine, EngineConfig};
 use crate::registry::model_registry::{CacheInfo, ModelInfo, ModelMetadata, ModelRegistry};
 
 /// Helper to create test app with a pre-registered test model
@@ -25,6 +25,7 @@ fn create_test_app() -> (axum::Router, TempDir) {
         MockEngine::new(),
         create_test_tokenizer(),
         "test-model".to_string(),
+        EngineConfig::default(),
     );
     tokio::spawn(runner.serve());
 

src/backend/llm_engine.rs

@@ -38,6 +38,8 @@ pub struct EngineHandle {
     event_tx: mpsc::UnboundedSender<SchedulerEvent>,
     seq_id_counter: Arc<AtomicU64>,
     model: String,
+    /// Completion-token budget to use when a request omits `max_tokens`.
+    default_max_tokens: usize,
 }
 
 impl EngineHandle {
@@ -129,6 +131,11 @@ impl EngineHandle {
     pub fn model(&self) -> &str {
         &self.model
     }
+
+    /// Completion-token budget to apply when a request omits `max_tokens`.
+    pub fn default_max_tokens(&self) -> usize {
+        self.default_max_tokens
+    }
 }
 
 /// The engine itself: owns the scheduler and backend, drives the event loop.
@@ -336,26 +343,80 @@ fn stream_flush(decoded: &str, sent_len: usize) -> Option<&str> {
     }
 }
 
+/// Tunable engine parameters.
+///
+/// These were previously hardcoded inside [`engine`]. Construct via
+/// [`EngineConfig::default`] and override fields as needed:
+///
+/// ```
+/// # use puma::backend::llm_engine::EngineConfig;
+/// let cfg = EngineConfig { max_batch_size: 64, ..Default::default() };
+/// ```
+#[derive(Debug, Clone)]
+pub struct EngineConfig {
+    /// Total KV-cache memory pool, in bytes.
+    pub memory_pool_bytes: usize,
+    /// Size of a single KV block, in bytes.
+    pub block_size_bytes: usize,
+    /// Number of tokens whose KV state fits in one block. Must be > 0.
+    pub tokens_per_block: usize,
+    /// Maximum number of sequences batched together per step.
+    pub max_batch_size: usize,
+    /// Completion-token budget applied when a request does not specify one.
+    pub default_max_tokens: usize,
+}
+
+impl EngineConfig {
+    /// Default KV-cache memory pool: 100 MB.
+    pub const DEFAULT_MEMORY_POOL_BYTES: usize = 1024 * 1024 * 100;
+    /// Default KV block size, in bytes.
+    pub const DEFAULT_BLOCK_SIZE_BYTES: usize = 512;
+    /// Default tokens per KV block.
+    pub const DEFAULT_TOKENS_PER_BLOCK: usize = 16;
+    /// Default maximum batch size.
+    pub const DEFAULT_MAX_BATCH_SIZE: usize = 32;
+    /// Default completion-token budget when a request omits `max_tokens`.
+    pub const DEFAULT_MAX_TOKENS: usize = 100;
+}
+
+impl Default for EngineConfig {
+    fn default() -> Self {
+        Self {
+            memory_pool_bytes: Self::DEFAULT_MEMORY_POOL_BYTES,
+            block_size_bytes: Self::DEFAULT_BLOCK_SIZE_BYTES,
+            tokens_per_block: Self::DEFAULT_TOKENS_PER_BLOCK,
+            max_batch_size: Self::DEFAULT_MAX_BATCH_SIZE,
+            default_max_tokens: Self::DEFAULT_MAX_TOKENS,
+        }
+    }
+}
+
 /// Construct a paired [`EngineHandle`] and [`EngineRunner`].
 ///
-/// Spawn `runner.serve()` on a task and share the returned handle with the API /
-/// CLI. All request submission goes through events, so the handle never
-/// touches the scheduler directly.
+/// Pass [`EngineConfig::default`] for the standard settings, or override fields
+/// to tune memory/batching. Spawn `runner.serve()` on a task and share the
+/// returned handle with the API / CLI. All request submission goes through
+/// events, so the handle never touches the scheduler directly.
 pub fn engine<B: Backend + Clone + 'static>(
     backend: B,
     tokenizer: Tokenizer,
     model: String,
+    config: EngineConfig,
 ) -> (EngineHandle, EngineRunner<B>) {
-    // Create block manager (100MB memory pool, 512 bytes per block)
-    let allocator = Box::new(CpuAllocator::new(1024 * 1024 * 100));
-    let block_manager = BlockManager::new(allocator, 512);
+    // KV-cache memory pool + block manager.
+    let allocator = Box::new(CpuAllocator::new(config.memory_pool_bytes));
+    let block_manager = BlockManager::new(allocator, config.block_size_bytes);
 
     // Event channel: the handle produces events, the scheduler consumes them.
     let (event_tx, event_rx) = mpsc::unbounded_channel();
 
-    // Create scheduler (max 32 batch size, 16 tokens per block); it owns the
-    // event receiver.
-    let scheduler = Scheduler::new(block_manager, event_rx, 32, 16);
+    // The scheduler owns the event receiver.
+    let scheduler = Scheduler::new(
+        block_manager,
+        event_rx,
+        config.max_batch_size,
+        config.tokens_per_block,
+    );
 
     let tokenizer = Arc::new(tokenizer);
 
@@ -364,6 +425,7 @@ pub fn engine<B: Backend + Clone + 'static>(
         event_tx,
         seq_id_counter: Arc::new(AtomicU64::new(1)),
         model,
+        default_max_tokens: config.default_max_tokens,
     };
 
     let runner = EngineRunner {
@@ -392,7 +454,12 @@ mod tests {
     async fn test_llm_engine() {
         let backend = MockEngine::new();
         let tokenizer = create_test_tokenizer();
-        let (handle, runner) = engine(backend, tokenizer, "test-model".to_string());
+        let (handle, runner) = engine(
+            backend,
+            tokenizer,
+            "test-model".to_string(),
+            EngineConfig::default(),
+        );
 
         tokio::spawn(runner.serve());
 

src/backend/mod.rs

@@ -2,4 +2,4 @@ pub mod engine;
 pub mod llm_engine;
 pub mod mock;
 
-pub use llm_engine::{engine, EngineHandle};
+pub use llm_engine::{engine, EngineConfig, EngineHandle};