From d0ccb02711d8352f23f43044c4f5c3ef75d66291 Mon Sep 17 00:00:00 2001 From: bluecoffee8 Date: Wed, 10 Jun 2026 17:23:20 -0700 Subject: [PATCH] add design docs --- python/minisgl/attention/DESIGN.md | 140 +++++++++++++++++++++++ python/minisgl/distributed/DESIGN.md | 111 ++++++++++++++++++ python/minisgl/engine/DESIGN.md | 136 ++++++++++++++++++++++ python/minisgl/kernel/DESIGN.md | 124 ++++++++++++++++++++ python/minisgl/kvcache/DESIGN.md | 138 ++++++++++++++++++++++ python/minisgl/layers/DESIGN.md | 142 +++++++++++++++++++++++ python/minisgl/message/DESIGN.md | 104 +++++++++++++++++ python/minisgl/models/DESIGN.md | 165 +++++++++++++++++++++++++++ python/minisgl/moe/DESIGN.md | 119 +++++++++++++++++++ python/minisgl/scheduler/DESIGN.md | 165 +++++++++++++++++++++++++++ python/minisgl/server/DESIGN.md | 137 ++++++++++++++++++++++ python/minisgl/tokenizer/DESIGN.md | 112 ++++++++++++++++++ 12 files changed, 1593 insertions(+) create mode 100644 python/minisgl/attention/DESIGN.md create mode 100644 python/minisgl/distributed/DESIGN.md create mode 100644 python/minisgl/engine/DESIGN.md create mode 100644 python/minisgl/kernel/DESIGN.md create mode 100644 python/minisgl/kvcache/DESIGN.md create mode 100644 python/minisgl/layers/DESIGN.md create mode 100644 python/minisgl/message/DESIGN.md create mode 100644 python/minisgl/models/DESIGN.md create mode 100644 python/minisgl/moe/DESIGN.md create mode 100644 python/minisgl/scheduler/DESIGN.md create mode 100644 python/minisgl/server/DESIGN.md create mode 100644 python/minisgl/tokenizer/DESIGN.md diff --git a/python/minisgl/attention/DESIGN.md b/python/minisgl/attention/DESIGN.md new file mode 100644 index 00000000..573731a1 --- /dev/null +++ b/python/minisgl/attention/DESIGN.md @@ -0,0 +1,140 @@ +# Attention Backends + +The `attention/` component provides pluggable attention kernels. All backends implement the same `BaseAttnBackend` interface so the rest of the system is kernel-agnostic. + +--- + +## Backend Hierarchy + +``` +BaseAttnBackend (abstract) + │ + ├── FlashAttentionBackend (fa) ← sgl_kernel flash_attn_with_kvcache + ├── FlashInferBackend (fi) ← flashinfer BatchPrefill/Decode wrappers + ├── TRTLLMBackend (trtllm) ← TensorRT-LLM paged attention + └── HybridBackend ← delegates to prefill_backend / decode_backend +``` + +**Auto-selection** at startup (when `attention_backend = "auto"`): +``` +SM 10.0+ (Blackwell) → "trtllm" +SM 9.0+ (Hopper) → "fa,fi" (FA for prefill, FlashInfer for decode) +otherwise → "fi" (FlashInfer only) +``` + +`"fa,fi"` creates a `HybridBackend(prefill=FA, decode=FI)`. + +--- + +## BaseAttnBackend Interface + +```python +forward(q, k, v, layer_id, batch) → Tensor # compute attention + store KV +prepare_metadata(batch) # build kernel-specific metadata before forward +init_capture_graph(max_seq_len, bs_list) # allocate static buffers for CUDA graph capture +prepare_for_capture(batch) # set up a specific bs for graph capture +prepare_for_replay(batch) # update static buffers before graph replay +``` + +--- + +## Data Flow per Layer + +``` +AttentionLayer.forward(qkv) + │ + ├── split qkv → q, k, v + ├── (optional) q_norm, k_norm + ├── rotary embedding + │ + ▼ +ctx.attn_backend.forward(q, k, v, layer_id, batch) + │ + ├── store_kv(k, v, batch.out_loc, layer_id) ← scatter K/V into paged cache + └── run kernel(q, k_cache, v_cache, metadata) → output +``` + +`batch.out_loc` is pre-computed by the scheduler: physical token addresses in the KV cache where this layer's K/V should be written. + +--- + +## FlashAttentionBackend (FA) + +Uses `sgl_kernel.flash_attn.flash_attn_with_kvcache`. + +**Metadata** (`FAMetadata`): +``` +cu_seqlens_q [bs+1] cumulative query sequence lengths +cu_seqlens_k [bs+1] cumulative key sequence lengths (including cached) +cache_seqlens [bs] total K length per request +page_table [bs, max_pages] page indices (divided by page_size) +max_seqlen_q / max_seqlen_k +``` + +Supports paged KV via page_table. FA version 3 (Hopper) or 4 (Blackwell) selected automatically. + +--- + +## FlashInferBackend (FI) + +Uses `flashinfer` `BatchPrefillWithPagedKVCacheWrapper` and `BatchDecodeWithPagedKVCacheWrapper`. + +**Metadata** (`FIMetadata`): +``` +cu_seqlens_q/k (CPU + GPU copies) +indices [total_tokens] flat list of physical page addresses (page_size=1) +last_page_len [bs] always 1 (since page_size=1) +wrapper ← prefill or decode wrapper, chosen per batch.is_prefill +``` + +FlashInfer requires `page_size=1` — the KV cache is treated as a flat token array. + +`_initialize_metadata_once`: FlashInfer's `.plan()` must be called once per metadata object before `.run()`. A CUDA event serializes plan calls to avoid buffer races. + +**Tensor cores** for decode: enabled when GQA ratio ≥ 4 (i.e., `num_qo_heads / num_kv_heads ≥ 4`). + +--- + +## CUDA Graph Integration + +Each backend has a two-phase graph protocol: + +``` +Capture phase (startup): + init_capture_graph(max_seq_len, bs_list) ← allocate static CUDAGraph buffers + for each bs: + prepare_for_capture(batch) ← point metadata at static buffers + torch.cuda.graph(...) + model.forward() ← recorded + +Replay phase (runtime): + prepare_for_replay(batch) ← copy live metadata into static buffers + graph.replay() +``` + +`HybridBackend` delegates `init_capture_graph`, `prepare_for_capture`, `prepare_for_replay` to `decode_backend` only (prefill is never graph-captured). + +--- + +## AttnMetadata.get_last_indices + +Used by the sampler to extract the logit for the last token of each request: + +``` +FA: cu_seqlens_q[1:bs+1] - 1 +FI: cu_seqlens_q_gpu[1:bs+1] - 1 +``` + +This gives the index of the last query token position within the flattened token batch. + +--- + +## Key Files + +| File | Responsibility | +|------|---------------| +| `base.py` | `BaseAttnBackend`, `BaseAttnMetadata`, `HybridBackend` | +| `fa.py` | `FlashAttentionBackend`, `FAMetadata` | +| `fi.py` | `FlashInferBackend`, `FIMetadata` | +| `trtllm.py` | `TRTLLMBackend` | +| `utils.py` | `BaseCaptureData` shared between backends | diff --git a/python/minisgl/distributed/DESIGN.md b/python/minisgl/distributed/DESIGN.md new file mode 100644 index 00000000..666f8044 --- /dev/null +++ b/python/minisgl/distributed/DESIGN.md @@ -0,0 +1,111 @@ +# Distributed + +The `distributed/` component handles tensor parallelism (TP) communication across GPU ranks. It provides a small abstraction layer that lets model layers call `all_reduce` / `all_gather` without knowing which underlying transport is used. + +--- + +## Component Map + +``` +DistributedCommunicator ← singleton used by model layers + └── plugins: List[DistributedImpl] ← stack; last one wins + +DistributedImpl (abstract) + ├── TorchDistributedImpl ← default; uses torch.distributed (NCCL or gloo) + └── PyNCCLDistributedImpl ← custom PyNCCL via CUDA IPC, faster for TP + └── PyNCCLCommunicator ← C++ NCCL wrapper (kernel/) +``` + +--- + +## Transport Selection + +``` +Engine.__init__ + │ + ├── tp_size == 1 OR use_pynccl flag? + │ ├── YES → init_process_group(backend="gloo") + │ │ enable_pynccl_distributed(...) + │ │ └── DistributedCommunicator.plugins.append(PyNCCLDistributedImpl) + │ │ + │ └── NO → init_process_group(backend="nccl") + │ (uses torch.distributed NCCL directly) + │ + └── tp_cpu_group = gloo group ← always used for CPU-side coordination + (free memory all-reduce, sync barriers) +``` + +When PyNCCL is active it **replaces** torch distributed for GPU-to-GPU transfers by being pushed onto the `plugins` stack. `all_reduce` / `all_gather` always call `plugins[-1]`. + +--- + +## Communication Primitives + +``` +DistributedCommunicator.all_reduce(x) + └── plugins[-1].all_reduce(x) + ├── TorchDistributedImpl → dist.all_reduce(x, SUM) + └── PyNCCLDistributedImpl → comm.all_reduce(x, "sum") + +DistributedCommunicator.all_gather(x) + └── plugins[-1].all_gather(x) + ├── TorchDistributedImpl → dist.all_gather_into_tensor(out, x) + └── PyNCCLDistributedImpl → comm.all_gather(result, x) + output shape: [world_size * x.shape[0], ...] +``` + +TP rank = 1: both implementations short-circuit and return `x` unchanged. + +--- + +## DistributedInfo + +```python +@dataclass +class DistributedInfo: + rank: int + size: int + + def is_primary(self) -> bool: + return self.rank == 0 +``` + +Used throughout the system to shard weights and KV heads across ranks. Accessed via `get_tp_info()` (process-local singleton set at engine init). + +--- + +## How Model Layers Use This + +``` +LinearOProj.forward(x): ← output projection (row-parallel) + y = F.linear(x, self.weight) + if tp_size > 1: + y = DistributedCommunicator().all_reduce(y) + return y + +VocabParallelEmbedding.forward: ← uses all_reduce after gather +LinearRowParallel.forward: ← same pattern as OProj +``` + +Column-parallel layers (QKV, gate/up projections) shard output dim; each rank computes a slice and no communication is needed until the row-parallel all-reduce. + +--- + +## Lifecycle + +``` +startup: enable_pynccl_distributed(...) ← append PyNCCL to plugins +shutdown: destroy_distributed() ← clear plugins list + torch.distributed.destroy_process_group() +``` + +`destroy_distributed` must be called **before** freeing NCCL resources to prevent hangs (ordering enforced in `Engine.shutdown` → `GraphRunner.destroy_cuda_graphs` first). + +--- + +## Key Files + +| File | Responsibility | +|------|---------------| +| `impl.py` | `DistributedImpl`, `TorchDistributedImpl`, `PyNCCLDistributedImpl`, `DistributedCommunicator` | +| `info.py` | `DistributedInfo`, `get_tp_info`, `set_tp_info` | diff --git a/python/minisgl/engine/DESIGN.md b/python/minisgl/engine/DESIGN.md new file mode 100644 index 00000000..cfb45311 --- /dev/null +++ b/python/minisgl/engine/DESIGN.md @@ -0,0 +1,136 @@ +# Engine + +The `engine/` component is the GPU-side execution core. It owns the model weights, KV cache allocation, CUDA graph replay, and token sampling. The scheduler calls into it once per batch. + +--- + +## Component Map + +``` +EngineConfig + │ + ▼ +Engine + ├── model (BaseLLMModel) ← loaded from HuggingFace weights + ├── kv_cache (BaseKVCachePool) ← paged GPU memory for K/V tensors + ├── page_table (int32 tensor) ← [max_req+1, max_seq_len] physical location lookup + ├── attn_backend ← FlashAttention / FlashInfer / TRTLLM + ├── moe_backend ← (optional) fused MoE kernel + ├── sampler (Sampler) ← greedy / top-k / top-p token selection + └── graph_runner (GraphRunner) ← CUDA graph capture + replay for decode +``` + +--- + +## Initialization Sequence + +``` +1. Set CUDA device, random seed, CUDA stream +2. Init distributed (gloo + optional pynccl) +3. Measure free GPU memory ─────────────────────────────┐ +4. Load model weights onto GPU │ used to compute +5. Measure free GPU memory again ───────────────────────┘ KV cache budget +6. Allocate KV cache pages (memory_ratio * init_free − model_memory) +7. Allocate page table tensor +8. Create attention backend +9. Create MoE backend (if MoE model) +10. Create Sampler +11. Capture CUDA graphs (GraphRunner) +``` + +--- + +## `forward_batch` Call Flow + +``` +Scheduler calls forward_batch(batch, args) + │ + ▼ + ctx.forward_batch(batch) ← sets global context so layers can read batch + │ + ├─ can_use_cuda_graph? + │ ├── YES → graph_runner.replay(batch) (decode, small bs) + │ └── NO → model.forward() (prefill or large bs) + │ + ▼ + req.complete_one() for each req ← advances cached_len / device_len + │ + ▼ + sampler.sample(logits, args) → next_tokens_gpu + │ + ▼ + async D2H copy of next_tokens_cpu + │ + ▼ + ForwardOutput(next_tokens_gpu, next_tokens_cpu, copy_done_event) +``` + +--- + +## Memory Budget Calculation + +``` +init_free_memory (before weight load) +model_memory = init_free - post_load_free +available_memory = memory_ratio × init_free - model_memory +num_pages = available_memory ÷ cache_per_page + +cache_per_page = 2 (K+V) × head_dim × local_kv_heads × page_size × dtype_bytes × num_layers +``` + +All TP ranks synchronize their free-memory values via CPU all-reduce and take the minimum so KV cache sizes are identical across ranks. + +--- + +## GraphRunner (CUDA Graph Capture) + +CUDA graphs pre-record the full decode forward pass for each supported batch size, eliminating CPU kernel-launch overhead during decode. + +``` +Startup (capture phase) +───────────────────── +For each bs in [1, 2, 4, 8, 16, ..., max_bs]: + create dummy Batch(reqs=[dummy_req] * bs) + warmup run (not recorded) + torch.cuda.graph(...) + └── model.forward() ← recorded into CUDAGraph + graph_map[bs] = graph + +Runtime (replay phase) +────────────────────── +batch arrives (decode, bs ≤ max_graph_bs) + copy batch tensors → GraphCaptureBuffer (pre-allocated static buffers) + attn_backend.prepare_for_replay(batch) ← update paged KV metadata + graph_map[padded_bs].replay() + return buffer.logits[:real_bs] +``` + +Batch sizes are padded to the next captured size using a dummy request that points to a dummy KV cache page, so no out-of-bounds reads occur. + +--- + +## Sampler + +``` +BatchSamplingArgs (per-batch, GPU) + ├── temperatures [bs] + ├── top_k [bs] + └── top_p [bs] + +sampler.sample(logits, args) + ├── temperature scaling + ├── top-k filtering + ├── top-p (nucleus) filtering + └── multinomial or argmax → token ids [bs] +``` + +--- + +## Key Files + +| File | Responsibility | +|------|---------------| +| `engine.py` | Engine class, init, `forward_batch` | +| `config.py` | `EngineConfig` dataclass | +| `graph.py` | `GraphRunner`, CUDA graph capture/replay | +| `sample.py` | `Sampler`, `BatchSamplingArgs` | diff --git a/python/minisgl/kernel/DESIGN.md b/python/minisgl/kernel/DESIGN.md new file mode 100644 index 00000000..2e3fb436 --- /dev/null +++ b/python/minisgl/kernel/DESIGN.md @@ -0,0 +1,124 @@ +# Kernel + +The `kernel/` component provides low-level GPU kernels and C++ extensions used throughout mini-sglang. It includes custom CUDA kernels compiled via PyTorch JIT, Triton kernels, and a Python-level NCCL wrapper. + +--- + +## Component Map + +``` +kernel/ +├── Python wrappers +│ ├── index.py ← fast_index_put, fast_compare_key +│ ├── store.py ← store_cache (scatter K/V into paged memory) +│ ├── radix.py ← fast_compare_key (used by RadixTreeNode) +│ ├── tensor.py ← Tensor C++ extension (serialization) +│ ├── pynccl.py ← PyNCCL communicator wrapper +│ └── utils.py ← fused_moe_kernel_triton, moe_sum_reduce_triton +│ +├── Triton kernel +│ └── triton/fused_moe.py ← Triton MoE matmul kernel +│ +└── CUDA C++ sources (csrc/) + ├── src/pynccl.cu ← custom NCCL all_reduce / all_gather + ├── src/radix.cpp ← fast byte-level key compare + ├── src/tensor.cpp ← tensor serialization for IPC + ├── jit/index.cu ← scatter-index kernels + └── jit/store.cu ← paged KV store kernel +``` + +--- + +## store_cache — KV Scatter Kernel + +``` +store_cache(k_cache, v_cache, indices, k, v) + +k_cache: [total_slots, heads, head_dim] ← flattened paged KV tensor +v_cache: [total_slots, heads, head_dim] +indices: [num_tokens] ← physical slot addresses +k, v: [num_tokens, heads, head_dim] ← new KV from forward pass + +Operation: k_cache[indices] = k + v_cache[indices] = v + +Implemented as a CUDA scatter write (index.cu / store.cu). +``` + +This is called once per transformer layer during forward, writing only the newly computed tokens (not the cached prefix). + +--- + +## fast_compare_key — Radix Tree Key Comparison + +``` +fast_compare_key(node_key: Tensor, input_ids: Tensor) → int + +Returns the length of the common prefix between node_key and input_ids. +Used by RadixTreeNode.get_match_len() during prefix cache tree walk. + +Implemented in C++ (radix.cpp) for speed; avoids Python loop overhead +on potentially long token sequences. +``` + +--- + +## PyNCCL — Custom NCCL Wrapper + +``` +init_pynccl(tp_rank, tp_size, tp_cpu_group, max_size_bytes) + → PyNCCLCommunicator + +PyNCCLCommunicator: + .all_reduce(tensor, op="sum") + .all_gather(output, input) +``` + +PyNCCL avoids the PyTorch distributed NCCL overhead (Python GIL, extra dispatching) by calling NCCL operations directly from C++ via CUDA IPC. Particularly beneficial for small all-reduce tensors in tensor parallelism (e.g., hidden states after each linear layer). + +A single pre-allocated GPU buffer of size `max_size_bytes` is reused for all NCCL operations to avoid dynamic allocation. + +--- + +## Triton MoE Kernel + +``` +fused_moe_kernel_triton( + hidden_states, w, + output, + topk_weights, topk_ids, + sorted_token_ids, expert_ids, num_tokens_post_padded, + apply_router_weight, + topk, config, + compute_type +) +``` + +A Triton kernel that performs a batched matrix multiplication across multiple experts simultaneously. Tokens are sorted by expert assignment, padded to `BLOCK_SIZE_M`, and processed in tiled groups. Block sizes are tuned via `config` (see `moe/DESIGN.md`). + +``` +moe_sum_reduce_triton(intermediate_cache3, out_hidden_states) + ← weighted sum of topk expert outputs back into hidden_states +``` + +--- + +## JIT Compilation + +CUDA kernels in `csrc/jit/` are compiled on first use via PyTorch's `torch.utils.cpp_extension.load`. Compiled artifacts are cached in the system's temp dir. The `kernel/__main__.py` script can be used to pre-compile all kernels ahead of time. + +--- + +## Key Files + +| File | Responsibility | +|------|---------------| +| `store.py` | `store_cache` wrapper | +| `index.py` | `fast_index_put` | +| `radix.py` | `fast_compare_key` | +| `pynccl.py` | `init_pynccl`, `PyNCCLCommunicator` | +| `utils.py` | `fused_moe_kernel_triton`, `moe_sum_reduce_triton` | +| `triton/fused_moe.py` | Triton MoE matmul kernel | +| `csrc/src/pynccl.cu` | NCCL CUDA C++ | +| `csrc/src/radix.cpp` | Fast key compare C++ | +| `csrc/jit/store.cu` | KV scatter CUDA kernel | diff --git a/python/minisgl/kvcache/DESIGN.md b/python/minisgl/kvcache/DESIGN.md new file mode 100644 index 00000000..9d28fee3 --- /dev/null +++ b/python/minisgl/kvcache/DESIGN.md @@ -0,0 +1,138 @@ +# KV Cache + +The `kvcache/` component provides two orthogonal abstractions: + +1. **KV Cache Pool** — GPU memory holding the actual key/value tensors, indexed by physical page. +2. **Prefix Cache** — a logical index mapping token sequences to physical page addresses, enabling prompt-prefix reuse across requests. + +--- + +## Abstraction Layers + +``` +┌───────────────────────────────────────────────────────────┐ +│ Scheduler / Engine │ +└────────────────────────────┬──────────────────────────────┘ + │ + ┌───────────────┴───────────────┐ + ▼ ▼ + BasePrefixCache BaseKVCachePool + (logical index) (physical GPU memory) + │ │ + ┌────┴─────┐ ┌────────┴────────┐ + │ Radix │ │ MHAKVCache │ + │ Cache │ │ (paged tensors) │ + └──────────┘ └─────────────────┘ + NaiveCache + (no reuse) +``` + +--- + +## BaseKVCachePool — Physical Storage + +`MHAKVCache` allocates a single contiguous tensor for all layers: + +``` +_kv_buffer: Tensor[2, num_layers, num_pages+1, page_size, local_kv_heads, head_dim] + │ + ├── [0] → K buffer (k_cache) + └── [1] → V buffer (v_cache) + + num_pages+1 includes one "dummy page" (index num_pages) + pointed to by dummy requests during CUDA graph capture. +``` + +**store_kv(k, v, out_loc, layer_id)** + +``` +out_loc: Tensor[total_tokens] ← physical token addresses + │ + ▼ +k_cache[layer_id].view(num_pages*page_size, heads, dim)[out_loc] ← k +v_cache[layer_id].view(num_pages*page_size, heads, dim)[out_loc] ← v +``` + +Writing uses the `store_cache` custom kernel for coalesced scatter writes. + +--- + +## BasePrefixCache — Logical Index + +### RadixPrefixCache + +A radix tree (compressed trie) keyed on page-aligned token sequences. Each tree node stores a slice of token IDs and their corresponding physical page addresses. + +``` +root (always protected, ref_count=1) + ├── [tok0..tokN] → node_A (pages [p0..pN]) + │ └── [tokN+1..tokM] → node_B (pages [pN+1..pM]) + └── [tok0..tokK] → node_C ... + +ref_count > 0 → "protected" (in use by an active request, cannot evict) +ref_count = 0 → "evictable" (LRU candidate) +``` + +**match_prefix(input_ids)** +- Walk tree greedily, aligning matches to `page_size` boundaries. +- Split a node if only a prefix of it matches. +- Returns a `RadixCacheHandle(cached_len, node)`. + +**insert_prefix(input_ids, indices)** +- Walk to deepest match, then append a new node for the unmatched suffix (page-aligned). +- Returns how many tokens were already in cache before insertion. + +**evict(size)** +- Collect leaf nodes with `ref_count == 0`, heap-sort by LRU timestamp. +- Pop leaves until enough pages freed; orphaned parents that become leaves are added to heap. + +### NaiveCache + +No prefix reuse. Every `match_prefix` returns 0 cached tokens. Useful for benchmarking or correctness testing. + +--- + +## Handle & Lock Protocol + +``` +match_prefix(input_ids) → MatchResult(cuda_handle) + │ + │ lock_handle(handle) ← protect from eviction during scheduling + │ + ▼ + req begins prefill + │ + │ insert_prefix(...) ← add new KV content to cache + │ unlock_handle(handle) ← allow eviction of old handle + ▼ + req finishes → cache_req(req, finished=True) + └── free tail pages that couldn't be inserted +``` + +Locking increments `ref_count` on all ancestor nodes; unlocking decrements. A node is evictable only when `ref_count == 0`. + +--- + +## Size Accounting + +``` +SizeInfo(evictable_size, protected_size) + total_size = evictable_size + protected_size (in tokens) + +CacheManager.available_size + = prefix_cache.evictable_size + + len(free_slots) * page_size +``` + +The CacheManager merges free physical slots with evictable prefix cache pages into a single "available" budget used by the prefill adder. + +--- + +## Key Files + +| File | Responsibility | +|------|---------------| +| `base.py` | `BaseKVCachePool`, `BasePrefixCache`, `BaseCacheHandle`, `SizeInfo` | +| `mha_pool.py` | `MHAKVCache` — paged GPU tensor storage | +| `radix_cache.py` | `RadixPrefixCache`, `RadixTreeNode`, LRU eviction | +| `naive_cache.py` | `NaivePrefixCache` — no-op prefix matching | diff --git a/python/minisgl/layers/DESIGN.md b/python/minisgl/layers/DESIGN.md new file mode 100644 index 00000000..22a6b703 --- /dev/null +++ b/python/minisgl/layers/DESIGN.md @@ -0,0 +1,142 @@ +# Layers + +The `layers/` component provides reusable tensor-parallel neural network building blocks. All layers are stateless operations (no `nn.Module`), use global context for batch info, and shard weights across TP ranks automatically. + +--- + +## Layer Taxonomy + +``` +BaseOP ← base class for all layers (forward() convention) +│ +├── StateLessOP ← no weight tensors (e.g., AttentionLayer, RoPE) +│ +├── Linear variants +│ ├── LinearReplicated ← full weight on every rank (e.g., small heads) +│ ├── LinearColParallelMerged ← output sharded across ranks (gate/up proj) +│ ├── LinearQKVMerged ← merged Q/K/V with TP-aware head splitting +│ ├── LinearOProj ← row-parallel + all_reduce (attn output proj) +│ └── LinearRowParallel ← row-parallel + all_reduce (MLP down proj) +│ +├── Embedding variants +│ ├── VocabParallelEmbedding ← vocab sharded across ranks, all_reduce +│ └── ParallelLMHead ← column-parallel vocab projection +│ +├── AttentionLayer ← split QKV, apply RoPE, call attn_backend +├── RMSNorm / RMSNormFused ← standard root mean square layer norm +├── Activation ← silu_and_mul, gelu_and_mul (fused) +├── MoELayer ← gate + fused expert routing +└── RotaryEmbedding ← RoPE: precomputed sin/cos tables +``` + +--- + +## Tensor Parallelism in Linear Layers + +Columns are split (ColumnParallel): each rank computes a slice of the output. +Rows are split (RowParallel): each rank computes a partial sum, then all-reduce. + +``` +Full linear: Y = X W^T (hidden_size → output_size) + +ColumnParallel (gate, up, Q projections): + local_W = W[rank * chunk : (rank+1) * chunk, :] + local_Y = X @ local_W^T ← no communication needed here + +RowParallel / OProj (down, O projections): + local_X = X[:, rank * chunk : (rank+1) * chunk] ← input is pre-sharded + partial_Y = local_X @ W^T + Y = all_reduce(partial_Y) ← sum across ranks + +QKV special case (LinearQKVMerged): + Q heads: sharded evenly across ranks + K/V heads: sharded (or replicated if num_kv_heads < tp_size) +``` + +--- + +## AttentionLayer + +``` +forward(qkv: Tensor) ← [total_tokens, (qo_dim + 2*kv_dim)] + │ + ├── split → q [*, qo_dim], k [*, kv_dim], v [*, kv_dim] + ├── optional q_norm(q), k_norm(k) (e.g. Qwen3) + ├── rotary.forward(positions, q, k) ← apply RoPE in-place + ├── reshape q → [*, num_qo_heads, head_dim] + └── ctx.attn_backend.forward(q, k, v, layer_id, ctx.batch) + └── returns output [*, num_qo_heads, head_dim] + → reshape → [*, qo_attn_dim] +``` + +--- + +## RMSNorm (Fused) + +``` +RMSNormFused.forward(x, residual=None) + │ + ├── if residual: x = x + residual (fused add-then-norm kernel) + └── x_norm = x / rms(x) * weight + returns (x_norm, x_as_new_residual) +``` + +The fused variant keeps a running residual for Llama-style pre-norm architectures, saving a separate add kernel. + +--- + +## Rotary Embedding + +``` +RotaryEmbedding.forward(positions, q, k) + │ + ├── cos, sin ← precomputed table[positions] + ├── apply_rotary_pos_emb(q, cos, sin) ← in-place + └── apply_rotary_pos_emb(k, cos, sin) +``` + +Tables are precomputed at init up to `max_position`. Supports `rope_scaling` variants (Llama3, YaRN, etc.) via factory function `get_rope(...)`. + +--- + +## OPList + +```python +class OPList(BaseOP): + op_list: List[BaseOP] + # models iterate op_list manually (e.g. for layer in self.layers.op_list) +``` + +Thin container used as a `list` of decoder layers. + +--- + +## MoE Layer + +``` +MoELayer.forward(hidden_states) + │ + ├── gate_proj(hidden_states) ← router logits [tokens, num_experts] + └── ctx.moe_backend.forward( + hidden_states, w1, w2, + gating_output, topk, renormalize + ) + └── returns output [tokens, hidden_size] +``` + +Expert weights `w1`, `w2` are stored on each rank (expert parallelism or full replication depending on config). + +--- + +## Key Files + +| File | Responsibility | +|------|---------------| +| `base.py` | `BaseOP`, `StateLessOP`, `OPList` | +| `linear.py` | All TP linear variants | +| `attention.py` | `AttentionLayer` | +| `embedding.py` | `VocabParallelEmbedding`, `ParallelLMHead` | +| `norm.py` | `RMSNorm`, `RMSNormFused` | +| `rotary.py` | `RotaryEmbedding`, `get_rope` factory | +| `activation.py` | `silu_and_mul`, `gelu_and_mul` | +| `moe.py` | `MoELayer` | diff --git a/python/minisgl/message/DESIGN.md b/python/minisgl/message/DESIGN.md new file mode 100644 index 00000000..e95cedc3 --- /dev/null +++ b/python/minisgl/message/DESIGN.md @@ -0,0 +1,104 @@ +# Message + +The `message/` component defines the message types that flow between the three process tiers: frontend (HTTP server), tokenizer, and scheduler backend. All messages are serialized over ZMQ sockets. + +--- + +## Message Hierarchy + +``` + ┌─ Frontend ──────────┐ + │ BaseFrontendMsg │ + │ BatchFrontendMsg │ + │ UserReply │ + └─────────────────────┘ + ▲ + backend → frontend + +BaseTokenizerMsg + BatchTokenizerMsg ← batch wrapper + TokenizeMsg ← frontend → tokenizer (raw text) + DetokenizeMsg ← backend → detokenizer (next token) + AbortMsg ← frontend → tokenizer (cancel request) + +BaseBackendMsg + BatchBackendMsg ← batch wrapper + UserMsg ← tokenizer → backend (token IDs) + AbortBackendMsg ← tokenizer → backend (cancel) + ExitMsg ← signal backend to shut down +``` + +--- + +## Message Flow by Path + +``` +User HTTP request + │ TokenizeMsg(uid, text, sampling_params) + ▼ +Tokenizer Worker + │ UserMsg(uid, input_ids: Tensor, sampling_params) + ▼ +Scheduler Backend + │ DetokenizeMsg(uid, next_token: int, finished: bool) + ▼ +Detokenizer Worker + │ UserReply(uid, incremental_output: str, finished: bool) + ▼ +Frontend Manager → HTTP response stream +``` + +Abort path: +``` +Frontend → AbortMsg(uid) → Tokenizer → AbortBackendMsg(uid) → Scheduler +``` + +--- + +## Batch Wrappers + +To reduce ZMQ round-trips, multiple messages may be wrapped in a single batch envelope: + +``` +BatchBackendMsg(data: List[BaseBackendMsg]) +BatchFrontendMsg(data: List[BaseFrontendMsg]) +BatchTokenizerMsg(data: List[BaseTokenizerMsg]) +``` + +Receivers always unwrap these before processing individual messages. + +--- + +## Serialization + +All message classes use a `serialize_type` / `deserialize_type` utility: + +```python +serialize_type(msg) → {"__type__": "UserMsg", "uid": 1, "input_ids": ..., ...} +deserialize_type(globals(), json) → UserMsg(...) +``` + +The `__type__` field is the class name, looked up in the module's `globals()` dict. Tensors are serialized as lists and reconstructed as `torch.int32` CPU tensors. + +--- + +## Key Message Fields + +| Message | Key Fields | +|---------|-----------| +| `TokenizeMsg` | `uid`, `text` (str or chat list), `sampling_params` | +| `UserMsg` | `uid`, `input_ids` (CPU int32 Tensor), `sampling_params` | +| `DetokenizeMsg` | `uid`, `next_token` (int), `finished` (bool) | +| `UserReply` | `uid`, `incremental_output` (str), `finished` (bool) | +| `AbortMsg` / `AbortBackendMsg` | `uid` | + +--- + +## Key Files + +| File | Responsibility | +|------|---------------| +| `backend.py` | `BaseBackendMsg`, `UserMsg`, `BatchBackendMsg`, `ExitMsg`, `AbortBackendMsg` | +| `frontend.py` | `BaseFrontendMsg`, `UserReply`, `BatchFrontendMsg` | +| `tokenizer.py` | `BaseTokenizerMsg`, `TokenizeMsg`, `DetokenizeMsg`, `AbortMsg`, `BatchTokenizerMsg` | +| `utils.py` | `serialize_type`, `deserialize_type` | diff --git a/python/minisgl/models/DESIGN.md b/python/minisgl/models/DESIGN.md new file mode 100644 index 00000000..41d04a4d --- /dev/null +++ b/python/minisgl/models/DESIGN.md @@ -0,0 +1,165 @@ +# Models + +The `models/` component contains transformer model implementations. Each model is composed of reusable `layers/` primitives and follows a shared pattern: global context supplies the active batch, so `forward()` takes no arguments at the top level. + +--- + +## Model Registry + +``` +ModelConfig.arch_name (e.g. "LlamaForCausalLM") + │ + ▼ + create_model(config) → instantiate registered class + │ + ▼ + load_weight(model_path, device) → weight iterator + model.load_state_dict(weights) +``` + +Registration happens via a decorator on each model class: + +```python +@register_model("LlamaForCausalLM") +class LlamaForCausalLM(BaseLLMModel): ... +``` + +--- + +## Supported Models + +| Class | Architecture | +|-------|-------------| +| `LlamaForCausalLM` | Llama 2/3 (dense) | +| `MistralForCausalLM` | Mistral (dense, sliding window) | +| `Qwen2ForCausalLM` | Qwen2 (dense) | +| `Qwen3ForCausalLM` | Qwen3 (dense, with Q/K norm) | +| `Qwen3MoeForCausalLM` | Qwen3-MoE (sparse MoE) | + +--- + +## Model Layer Stack (Llama as canonical example) + +``` +LlamaForCausalLM.forward() + │ + ▼ +LlamaModel.forward(input_ids) ← from ctx.batch.input_ids + │ + ├── embed_tokens(input_ids) ← VocabParallelEmbedding + │ → x [total_tokens, hidden_size] + │ + ├── for layer in layers: + │ LlamaDecoderLayer.forward(x, residual) + │ ├── input_layernorm(x, residual) ← RMSNormFused (fused residual add) + │ ├── self_attn.forward(x) ← QKV proj → Attention → O proj + │ ├── post_attention_layernorm(x, residual) + │ └── mlp.forward(x) ← GatedMLP (gate+up → silu → down) + │ + └── final norm(x, residual) + → hidden [total_tokens, hidden_size] + +LlamaForCausalLM: + lm_head(hidden) ← ParallelLMHead + → logits [total_tokens, vocab_size] +``` + +--- + +## BaseLLMModel + +```python +class BaseLLMModel(BaseOP): + def forward(self) -> torch.Tensor: + # reads batch from get_global_ctx().batch + ... + + def load_state_dict(self, state_dict): + # recursively matches weight names to BaseOP children + # handles TP-sharded weights (split along correct dim) +``` + +Weight loading slices each tensor according to TP rank, so every GPU only holds its shard. + +--- + +## ModelConfig + +```python +@dataclass +class ModelConfig: + arch_name: str # model class name + hidden_size: int + num_layers: int + num_qo_heads: int + num_kv_heads: int + head_dim: int + vocab_size: int + num_experts: int # 0 for dense models + is_moe: bool + rms_norm_eps: float + tie_word_embeddings: bool + rotary_config: RotaryConfig + ... +``` + +Loaded from `config.json` in the model directory via `load_model_config(model_path)`. + +--- + +## Attention Variants + +``` +RopeAttn (most models) + QKVProj (LinearQKVMerged) → AttentionLayer(+ RoPE) → OProj (LinearOProj) + +Qwen3 adds Q/K per-head norm (q_norm, k_norm) inside AttentionLayer. +``` + +--- + +## MLP Variants + +``` +GatedMLP (Llama, Qwen, Mistral) + gate_proj (LinearColParallelMerged, outputs gate + up merged) + silu_and_mul(gate, up) + down_proj (LinearRowParallel) + +MoELayer (Qwen3-MoE) + gate (LinearReplicated → router logits) + moe_backend.forward(hidden, w1, w2, gating_output, topk, ...) +``` + +--- + +## Weight Loading + +``` +load_weight(model_path, device) → Iterator[(name, tensor)] + +Supports: + - .safetensors shards (HuggingFace standard) + - .bin / .pt files + +Weight names are remapped via per-model rename rules to match the +internal attribute hierarchy (e.g., "model.layers.0.self_attn.q_proj.weight" +→ the correct LinearQKVMerged slice). +``` + +--- + +## Key Files + +| File | Responsibility | +|------|---------------| +| `base.py` | `BaseLLMModel` | +| `config.py` | `ModelConfig`, `RotaryConfig`, `load_model_config` | +| `register.py` | `@register_model`, `create_model` | +| `weight.py` | `load_weight`, safetensors reader | +| `llama.py` | Llama decoder stack | +| `mistral.py` | Mistral (sliding-window variant) | +| `qwen2.py` | Qwen2 | +| `qwen3.py` | Qwen3 (Q/K norm) | +| `qwen3_moe.py` | Qwen3-MoE (sparse MoE) | +| `utils.py` | Shared `GatedMLP`, `RopeAttn` building blocks | diff --git a/python/minisgl/moe/DESIGN.md b/python/minisgl/moe/DESIGN.md new file mode 100644 index 00000000..400939f8 --- /dev/null +++ b/python/minisgl/moe/DESIGN.md @@ -0,0 +1,119 @@ +# MoE (Mixture of Experts) + +The `moe/` component provides pluggable Mixture-of-Experts inference backends. The model layer calls into the active backend without knowing which kernel implementation is used. + +--- + +## Component Map + +``` +BaseMoeBackend (abstract) + │ + └── FusedMoe ← Triton-based fused expert kernel + +Selected via config: moe_backend = "fused" (auto-selected for MoE models) +``` + +--- + +## Interface + +```python +class BaseMoeBackend: + def forward( + self, + hidden_states: Tensor, # [num_tokens, hidden_size] + w1: Tensor, # [num_experts, ffn_hidden*2, hidden_size] (gate+up merged) + w2: Tensor, # [num_experts, hidden_size, ffn_hidden] (down proj) + gating_output: Tensor, # [num_tokens, num_experts] (router logits) + topk: int, # number of experts per token + renormalize: bool, # normalize expert weights to sum=1 + activation: str, # "silu" or "gelu" + apply_router_weight_on_input: bool, + ) -> Tensor: # [num_tokens, hidden_size] +``` + +--- + +## FusedMoe Execution Flow + +``` +gating_output [tokens, E] + │ + ▼ +fused_topk(gating_output, topk, renormalize) + │ + ├── topk_softmax kernel (sgl_kernel) + └── topk_weights [tokens, topk], topk_ids [tokens, topk] + │ + ▼ +moe_align_block_size(topk_ids, BLOCK_SIZE_M, E) + │ + ├── sorts token→expert assignments + ├── pads each expert's token count to block_size + └── sorted_token_ids, expert_ids, num_tokens_post_padded + │ + ▼ +fused_moe_kernel_triton (pass 1) ← w1 matmul + hidden_states × w1[experts] + → intermediate_cache1 [tokens, topk, ffn_hidden*2] + │ + ▼ +silu_and_mul / gelu_and_mul ← gated activation + → intermediate_cache2 [tokens*topk, ffn_hidden] + │ + ▼ +fused_moe_kernel_triton (pass 2) ← w2 matmul + intermediate_cache2 × w2[experts] + → intermediate_cache3 [tokens, topk, hidden_size] + │ + ▼ +moe_sum_reduce_triton ← weighted sum over topk experts + → output [tokens, hidden_size] +``` + +--- + +## Memory Layout + +``` +Intermediate buffers (reusing a single cache tensor): + cache [tokens * topk * max(ffn_hidden*2, hidden_size)] + + intermediate_cache1 = cache[:tokens*topk*ffn_hidden*2].view(tokens, topk, ffn_hidden*2) + intermediate_cache3 = cache[:tokens*topk*hidden_size].view(tokens, topk, hidden_size) + intermediate_cache2 = torch.empty(tokens*topk, ffn_hidden) ← separate (after activation) +``` + +Reusing `cache` for both cache1 and cache3 (via different view sizes) halves peak intermediate memory. + +--- + +## Block Size Tuning + +``` +try_get_optimal_moe_config(w1_shape, w2_shape, topk, M=num_tokens) + │ + ├── M > E → BLOCK_SIZE_M=64, N=64, K=32, GROUP_M=8 (large batch) + └── M ≤ E → BLOCK_SIZE_M=16, N=32, K=64, GROUP_M=1 (small batch / decode) +``` + +Block sizes control Triton tiling for the expert matmuls. Small batch sizes (typical during decode) use smaller M blocks to avoid wasted compute. + +--- + +## Router Weight Application + +`apply_router_weight_on_input=True`: multiply by expert weight before the first matmul (fused into pass 1). +`apply_router_weight_on_input=False`: multiply in the final reduce (pass 2, default). + +Both modes produce identical outputs; the flag allows fusing the weight into whichever pass is cheaper. + +--- + +## Key Files + +| File | Responsibility | +|------|---------------| +| `base.py` | `BaseMoeBackend` abstract interface | +| `fused.py` | `FusedMoe`, `fused_topk`, `moe_align_block_size`, `fused_experts_impl` | diff --git a/python/minisgl/scheduler/DESIGN.md b/python/minisgl/scheduler/DESIGN.md new file mode 100644 index 00000000..a0e0b481 --- /dev/null +++ b/python/minisgl/scheduler/DESIGN.md @@ -0,0 +1,165 @@ +# Scheduler + +The `scheduler/` component is the CPU-side brain of the inference loop. It runs in its own process per TP rank, orchestrates prefill/decode batching, manages KV cache page allocation, and drives the engine with prepared batches. + +--- + +## Component Map + +``` +Scheduler + ├── engine (Engine) ← GPU execution core + ├── table_manager (TableManager) ← allocates request "slots" (table indices) + ├── cache_manager (CacheManager) ← paged KV memory + prefix cache eviction + ├── prefill_manager (PrefillManager)← queue of pending / chunked prefill reqs + └── decode_manager (DecodeManager) ← set of in-flight decode reqs +``` + +--- + +## Main Loop: Overlap Scheduling + +The scheduler overlaps CPU metadata work for the *next* batch with GPU execution of the *current* batch. + +``` + ┌─────────────┐ ┌─────────────┐ + CPU thread │ receive msgs│ │ process last│ + │ schedule_ │ │ batch output│ + │ next_batch │ │ (tokens, │ + │ prepare_ │ │ detokenize,│ + │ batch │ │ free pages)│ + └──────┬──────┘ └──────┬──────┘ + │ launch │ + ┌──────▼──────────────────────▼──────┐ + GPU stream │ forward_batch(current) │ + └─────────────────────────────────────┘ + +Iteration N: schedule N → launch N → process N-1 result +Iteration N+1: schedule N+1 → launch N+1 → process N result +``` + +A separate CUDA stream is used for CPU-side scheduling metadata; engine runs on its own stream. `engine_stream.wait_stream(scheduler_stream)` ensures ordering. + +--- + +## Scheduling Priority + +``` +_schedule_next_batch() + │ + ├── prefill_manager.schedule_next_batch(prefill_budget) ← FIRST priority + │ └── returns None if no pending reqs + │ + └── decode_manager.schedule_next_batch() ← SECOND priority + └── returns None if no in-flight decode reqs +``` + +Prefill-first policy: ensures new requests make progress and enter the decode queue promptly. + +--- + +## Prefill Manager + +``` +pending_list: [PendingReq, ...] ← FIFO queue of new requests + +schedule_next_batch(prefill_budget): + PrefillAdder (token_budget=prefill_budget, reserved_size=decode_inflight_tokens) + for each pending_req: + ├── allocate table slot + ├── match prefix cache → cached_len (may skip tokens) + ├── chunk if extend_len > token_budget → ChunkedReq + └── add to batch + + Chunked reqs stay at the front of pending_list for next iteration. +``` + +**Chunked prefill**: if a request's input is too long to fit the token budget in one step, it is split across multiple prefill batches. `ChunkedReq` marks that the req must not be sampled yet. + +--- + +## Decode Manager + +``` +running_reqs: Set[Req] + +filter_reqs(reqs): ← called after each forward; adds newly-promoted reqs + running_reqs = {r for r in running_reqs ∪ reqs if r.can_decode} + +schedule_next_batch(): + Batch(reqs=sorted(running_reqs, key=uid), phase="decode") +``` + +Sorting by UID ensures stable ordering across TP ranks (critical for consistent sampling). + +--- + +## Cache Manager + +Sits between the prefix cache and the page table. Manages a free-slot pool of page-aligned slots. + +``` +free_slots: Tensor[int32] ← available page start addresses + +allocate_paged(reqs): + needed_pages ← sum of new pages each req needs + if needed_pages > free_slots: + evict from prefix_cache → reclaim pages + write newly allocated page addresses into page_table + +cache_req(req, finished): + insert req's input_ids + page indices → prefix_cache + free previously-matched pages (they're now in cache) + if finished: free the tail pages too + +lazy_free_region(): ← context manager + defers page frees until the GPU work using those pages completes +``` + +--- + +## Batch Preparation Pipeline + +``` +_prepare_batch(batch) + │ + ├── graph_runner.pad_batch(batch) ← pad to next CUDA graph size + ├── cache_manager.allocate_paged(reqs) ← write new pages to page_table + ├── _make_positions(batch) ← [total_tokens] position indices + ├── _make_input_tuple(batch) ← (table_idx, positions) for token lookup + ├── _make_write_tuple(batch) ← (table_idx, seq_len/-1) for KV write + ├── batch.out_loc ← page_table[input_mapping] + └── attn_backend.prepare_metadata(batch) + +Returns ForwardInput(batch, sample_args, input_tuple, write_tuple) +``` + +--- + +## Message Handling + +``` +receive_msg(blocking) + │ + ├── UserMsg → prefill_manager.add_one_req(msg) + ├── AbortBackendMsg → remove from prefill or decode manager, free resources + ├── ExitMsg → raise KeyboardInterrupt + └── BatchBackendMsg → unwrap and process each sub-message +``` + +After each forward pass, `_process_last_data` iterates tokens, sends `DetokenizeMsg` replies, and frees finished request resources (table slot, cache pages). + +--- + +## Key Files + +| File | Responsibility | +|------|---------------| +| `scheduler.py` | Main loop, overlap scheduling, message dispatch | +| `prefill.py` | `PrefillManager`, `PrefillAdder`, chunked prefill | +| `decode.py` | `DecodeManager`, decode batch scheduling | +| `cache.py` | `CacheManager`, page allocation/eviction | +| `table.py` | `TableManager`, request slot allocation | +| `io.py` | ZMQ I/O mixin (send/receive messages) | +| `config.py` | `SchedulerConfig` dataclass | +| `utils.py` | `PendingReq` helper | diff --git a/python/minisgl/server/DESIGN.md b/python/minisgl/server/DESIGN.md new file mode 100644 index 00000000..fa52d188 --- /dev/null +++ b/python/minisgl/server/DESIGN.md @@ -0,0 +1,137 @@ +# Server + +The `server/` component handles the HTTP API, process orchestration, and the IPC wiring between the frontend and backend workers. + +--- + +## Process Architecture + +``` +┌─────────────────────────────────────────────────────────────────────┐ +│ Main Process (FastAPI / uvicorn) │ +│ │ +│ FrontendManager │ +│ ├── ZmqAsyncPushQueue → zmq_tokenizer_addr │ +│ └── ZmqAsyncPullQueue ← zmq_frontend_addr │ +└──────────────────────────────────────────────────────────────────────┘ + │ ▲ + TokenizeMsg / AbortMsg UserReply + │ │ + ▼ │ +┌─────────────────────────────────────────────────────────────────────┐ +│ Tokenizer Processes (N = num_tokenizer) │ +│ tokenize_worker() │ +│ TokenizeMsg → input_ids → UserMsg → backend │ +│ DetokenizeMsg ← next_token ← backend → UserReply → frontend │ +│ AbortMsg → AbortBackendMsg → backend │ +└──────────────────────────────────────────────────────────────────────┘ + │ + UserMsg / AbortBackendMsg + │ + ▼ +┌─────────────────────────────────────────────────────────────────────┐ +│ Scheduler Processes (one per TP rank) │ +│ Scheduler.run_forever() │ +│ ├── GPU 0 (primary) → sends DetokenizeMsg back │ +│ └── GPU 1..N-1 → silent TP workers │ +└──────────────────────────────────────────────────────────────────────┘ +``` + +--- + +## Startup Sequence (`launch_server`) + +``` +1. parse_args() +2. run_api_server(config, start_backend_fn, run_shell) + │ + ├── create FrontendManager (ZMQ sockets) + ├── call start_backend_fn() ← spawns subprocesses: + │ ├── for i in range(tp_size): mp.Process(_run_scheduler) + │ ├── 1× detokenizer process mp.Process(tokenize_worker) + │ └── N× tokenizer processes mp.Process(tokenize_worker) + │ └── wait for ack_queue messages from all workers + │ + └── start uvicorn (or run interactive shell) +``` + +All subprocesses are spawned with `mp.set_start_method("spawn")`. They communicate via ZMQ IPC sockets (file paths configured in `ServerArgs`). + +--- + +## HTTP Endpoints + +| Method | Path | Description | +|--------|------|-------------| +| POST | `/generate` | Raw streaming generation (SSE, token-per-line) | +| POST | `/v1/chat/completions` | OpenAI-compatible chat completions (stream or non-stream) | +| GET | `/v1/models` | List available model(s) | +| GET/POST | `/v1` | Health check | + +--- + +## Request Lifecycle + +``` +HTTP POST /v1/chat/completions + │ + ├── uid = state.new_user() ← assign unique ID, create ack_map entry + ├── await state.send_one(TokenizeMsg(uid, text, sampling_params)) + │ + └── if stream: + return StreamingResponse(state.stream_chat_completions(uid)) + else: + async for ack in state.wait_for_ack(uid): + full_content += ack.incremental_output + return JSON response + +stream_chat_completions(uid): + async for ack in wait_for_ack(uid): + yield SSE chunk (OpenAI delta format) + yield final chunk (finish_reason="stop") + yield "data: [DONE]" +``` + +**Client disconnect detection**: `stream_with_cancellation` checks `request.is_disconnected()` on each chunk and sends `AbortMsg` if disconnected. + +--- + +## FrontendManager + +```python +@dataclass +class FrontendManager: + uid_counter: int ← monotonically increasing request ID + ack_map: Dict[int, List[UserReply]] ← buffered replies per uid + event_map: Dict[int, asyncio.Event] ← notifies when new reply arrives + +listen(): ← single asyncio task, drains recv queue + while True: + msg = await recv_tokenizer.get() + for reply in _unwrap_msg(msg): + ack_map[reply.uid].append(reply) + event_map[reply.uid].set() + +wait_for_ack(uid): ← async generator yielding UserReplys + while True: + await event.wait(); event.clear() + for ack in pending: yield ack + if ack.finished: break + cleanup ack_map / event_map +``` + +--- + +## Interactive Shell Mode + +When `--shell` flag is set, `run_shell=True` is passed to `run_api_server`, which calls `asyncio.run(shell())` instead of starting uvicorn. The shell is a `prompt_toolkit` REPL that maintains conversation history and calls the server's internal `shell_completion()` function directly. + +--- + +## Key Files + +| File | Responsibility | +|------|---------------| +| `launch.py` | `launch_server`, subprocess spawning, ack synchronization | +| `api_server.py` | FastAPI app, `FrontendManager`, HTTP endpoints, shell | +| `args.py` | `ServerArgs` — all configuration (host, port, ZMQ addrs, TP, model path) | diff --git a/python/minisgl/tokenizer/DESIGN.md b/python/minisgl/tokenizer/DESIGN.md new file mode 100644 index 00000000..096f8360 --- /dev/null +++ b/python/minisgl/tokenizer/DESIGN.md @@ -0,0 +1,112 @@ +# Tokenizer + +The `tokenizer/` component runs in a dedicated subprocess (or multiple) bridging the HTTP frontend and the GPU scheduler backend. It handles tokenization of incoming prompts and detokenization of generated token IDs into text. + +--- + +## Role in the System + +``` +Frontend (async HTTP) + │ TokenizeMsg / AbortMsg + ▼ +tokenize_worker (subprocess) + ├── tokenize → UserMsg → Scheduler backend + ├── detokenize ← DetokenizeMsg ← Scheduler backend + └── AbortMsg → AbortBackendMsg → Scheduler backend + │ UserReply + ▼ +Frontend (async HTTP) +``` + +There are two kinds of workers using the same `tokenize_worker` function: +- **Tokenizer workers** (`N = num_tokenizer`): receive from `zmq_tokenizer_addr`, handle `TokenizeMsg`. +- **Detokenizer worker** (`1`): receives from `zmq_detokenizer_addr`, handles `DetokenizeMsg`. + +Both send `UserMsg`/`AbortBackendMsg` to `zmq_backend_addr` and `UserReply` to `zmq_frontend_addr`. + +--- + +## tokenize_worker Loop + +```python +while True: + pending = [recv_listener.get()] ← blocking get + while len(pending) < local_bs and not recv_listener.empty(): + pending.extend(recv_listener.get()) ← opportunistic batching + + detokenize_msg = [m for m in pending if isinstance(m, DetokenizeMsg)] + tokenize_msg = [m for m in pending if isinstance(m, TokenizeMsg)] + abort_msg = [m for m in pending if isinstance(m, AbortMsg)] + + if detokenize_msg: + replies = DetokenizeManager.detokenize(detokenize_msg) + send UserReply(uid, incremental_text, finished) → frontend + + if tokenize_msg: + tensors = TokenizeManager.tokenize(tokenize_msg) + send UserMsg(uid, input_ids, sampling_params) → backend + + if abort_msg: + send AbortBackendMsg(uid) → backend +``` + +Opportunistic batching: after the first blocking receive, the worker drains any additional queued messages up to `local_bs` without blocking, processing them in one batch. + +--- + +## TokenizeManager + +``` +TokenizeManager.tokenize(msgs: List[TokenizeMsg]) + │ + ├── for each msg: + │ if msg.text is str → tokenizer.encode(text) + │ if msg.text is list (chat messages) → + │ tokenizer.apply_chat_template(messages, add_generation_prompt=True) + │ + └── returns List[Tensor[int32]] ← CPU tensors, one per request +``` + +Supports both raw string prompts and OpenAI-style message lists (role/content dicts). + +--- + +## DetokenizeManager + +``` +DetokenizeManager.detokenize(msgs: List[DetokenizeMsg]) + │ + ├── maintains per-uid state: List[int] (accumulated token ids) + ├── for each msg: + │ append msg.next_token to uid's buffer + │ decode(buffer) - decode(buffer[:-1]) ← incremental decode + │ (avoids multi-byte UTF-8 boundary issues) + │ + └── returns List[str] ← incremental text strings + +Finished requests: uid state is cleaned up when msg.finished == True +``` + +Incremental detokenization: re-decoding the full buffer minus the last token, then subtracting the previous string, correctly handles tokens that span byte boundaries (e.g., multi-byte UTF-8 characters). + +--- + +## Message Types Handled + +| Incoming | Direction | Action | +|----------|-----------|--------| +| `TokenizeMsg(uid, text, sampling_params)` | frontend → tokenizer | encode text → `UserMsg` | +| `DetokenizeMsg(uid, next_token, finished)` | backend → detokenizer | decode token → `UserReply` | +| `AbortMsg(uid)` | frontend → tokenizer | forward as `AbortBackendMsg` | +| `BatchTokenizerMsg` | any | unwrap and process each sub-message | + +--- + +## Key Files + +| File | Responsibility | +|------|---------------| +| `server.py` | `tokenize_worker` — main event loop, message routing | +| `tokenize.py` | `TokenizeManager` — text → token IDs | +| `detokenize.py` | `DetokenizeManager` — token ID → incremental text |