sgl-project · bluecoffee8 · Jun 11, 2026
diff --git a/python/minisgl/attention/DESIGN.md 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
@@ -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
@@ -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` |