gavin-d26/corellm

1. What this repo is

CoreLLM is a minimal LLM inference engine, not a training stack. Its architecture is centered on fast serving of a single supported family, Llama 3.1, with an emphasis on continuous batching, paged KV cache management, and multi-GPU tensor parallelism.

2. Top-level runtime modes

There are two entry paths:

- Offline Python API via EngineNode, used for batch inference directly from code.

- OpenAI-compatible server via corellm serve, which exposes /v1/completions and /v1/chat/completions for online requests.

Both paths converge on the same engine and worker runtime.

3. Request ingestion and client boundary

The server-side boundary is in corellm/openai_server. It accepts HTTP requests, maps them to OpenAI-style schemas, tokenizes prompts, and forwards work to the engine client. The client layer is the boundary between API concerns and model execution; it keeps tokenization/detokenization and request shaping off the GPU hot path.

4. Main engine control plane

corellm/engine/engine.py is the orchestrator. It owns:

- the serving loop,

- scheduler decisions,

- sequence state,

- KV-cache bookkeeping,

- and coordination with GPU workers.

The engine is the “control plane” of the system; it does not itself contain the model math.

5. Worker / model execution plane

GPU workers run model forward passes and own the model shards. The repo’s architecture description indicates a controller/worker split with:

- ZMQ for control messages,

- NCCL for tensor-parallel collectives.

This keeps HTTP, scheduling, and token management out of the worker process. The worker side is responsible for the actual forward compute and returns outputs/results to the engine.

6. Scheduling and continuous batching

corellm/engine/scheduler.py determines the next batch composition. The system is built around a single mixed forward pass that can contain both:

- prefill tokens from new prompts, and

- decode tokens from in-flight sequences.

The scheduler is decode-prioritized and chunks prefill to respect batching limits. This is the core throughput-oriented control loop in the architecture.

7. Sequence and KV-cache state

corellm/engine/sequence.py and corellm/engine/kv_cache/ hold per-request state and KV allocation logic. There are two cache allocators:

- naive

- radix-tree-based prefix caching

The KV cache is a major subsystem because it enables reuse of prompt prefixes and is tightly coupled to batch scheduling and memory layout.

8. Model graph and weight loading

corellm/engine/model_graph/ is the bridge from model definition to executable graph:

- executor.py runs the model graph on the worker side,

- weight_loader.py materializes pretrained weights into the model layout.

This layer separates model structure from runtime execution and is where checkpoint loading meets the inference graph.

9. Model implementation and reusable layers

corellm/models/ defines supported model wrappers and templates, currently centered on Llama 3.1. corellm/layers/ contains the reusable neural network primitives: attention, linear, embeddings, normalization, and activations. These are the semantic model components that the executor assembles into a forward pass.

10. Custom GPU kernels

corellm/kernels/ and corellm/engine/kv_cache/ are where the performance-critical CUDA/Triton pieces live. The repo explicitly includes:

- paged attention

- gather

- device-specific kernel configs

The README emphasizes a custom paged-attention kernel that handles mixed prefill/decode batches and arbitrary sequence lengths. This is the low-level compute boundary beneath the model graph.

11. Memory layout and parallelism

A key architectural decision is the single flat GPU allocation per rank: model weights occupy the front of the slab, and the remainder is carved into KV blocks. The amount of cache is derived from available device memory. On multi-GPU setups, the model is tensor-parallel, with sharded attention and fused column/row-parallel linears. This makes memory placement and collectives first-class architectural concerns.

12. Serving and benchmarking as separate subsystems

corellm/openai_server/ provides the public API surface, while benchmarks/ contains offline throughput, latency, serving, and kernel microbenchmarks. These benchmarking scripts are not part of the serving hot path, but they document the intended usage patterns and performance envelope of the engine. The results/ directory stores recorded benchmark outputs.

13. Tests and validation boundaries

The tests focus on integration behavior rather than unit-level numerical checks:

- tests/test_chat_batch.py

- tests/test_server_requests.py

This suggests the most important correctness boundaries are request batching and server/API behavior, not isolated math kernels.

14. Overall data flow to draw in a diagram

A request typically flows as:

HTTP/OpenAI request → server schemas → tokenization/client → engine scheduler/sequence state → worker forward pass → custom kernels on GPU → sampled tokens/results → detokenization → response streaming or return.

In offline mode, the same flow exists but without the HTTP/server front end. The main architectural boundary is between the control plane in the engine/server processes and the compute plane in GPU worker processes.