SGLang Notes (Mainly Awesome-ML-SYS-Tutorial Notes)

• Notes mainly from: https://github.com/zhaochenyang20/Awesome-ML-SYS-Tutorial/blob/main/README-cn.md
• Plus related topics explored while reading

Radix Attention

• Fairly straightforward; reference: https://github.com/CalvinXKY/InfraTech/blob/main/llm_infer/sglang_radix_attention.ipynb
• Attention with a KV cache
• Compress prefixes with a radix tree, similar to a trie
• Reference counting manages lifetimes

SGLang

• https://github.com/zhaochenyang20/Awesome-ML-SYS-Tutorial/blob/main/sglang/code-walk-through/readme-CN.md
  • Start http_server/grpc_server, containing TokenizerManager, DeTokenizerManager, and Scheduler. Tokenizer and Detokenizer mainly route messages; model logic lives in Scheduler. Its threading model uses an asynchronous event loop: requests enter a queue, then are collected into batches for model workers;
  • Supports PD separation with different servers for prefill and decode.
    • How do P and D interact? A prefill node calls _register_to_bootstrap to register its information with a bootstrap server, then starts its event loop. After processing produces tokens, it sends them to decode. The sending policy depends on disagg_kv_sender and varies by backend. For Mooncake, MooncakeKVManager has three roles and many KV-cache status checks coordinating cache turnover; pay particular attention to its thread/process model:
      1. prefill: execute requests and call disagg_kv_sender.send_kvcache();
      2. decode: connect to registered prefill nodes and poll for data. Two queues handle preallocation reservations and polling for completed KV transfers (poll())
      3. the transfer worker in prefill: primarily moves data, including shard mapping between P and D nodes
  • Supports overlap: sampling for the previous batch overlaps the next batch’s forward pass to reduce GPU waiting. The code is somewhat difficult; focus on this section:

1
@torch.no_grad()
2
def event_loop_overlap_disagg_prefill(self: Scheduler) -> None:
3
    self.result_queue = deque()
4

5
    while True:
6
        # Receive requests
7
        recv_reqs = self.recv_requests()
8
        self.process_input_requests(recv_reqs)
9
        self.waiting_queue.extend(
10
            self.disagg_prefill_bootstrap_queue.pop_bootstrapped()
11
        )
12

13
        # Get the next batch to run
14
        batch = self.get_next_disagg_prefill_batch_to_run()
15
        self.cur_batch = batch
16

17
        # Launch the current batch
18
        if batch:
19
            batch_result = self.run_batch(batch)
20
            self.result_queue.append((batch.copy(), batch_result))
21
        else:
22
            batch_result = None
23

24
        # Process the last batch
25
        if self.last_batch:
26
            tmp_batch, tmp_result = self.result_queue.popleft()
27
            self.process_batch_result_disagg_prefill(tmp_batch, tmp_result)
28
        elif batch is None:
29
            # When the server is idle, do self-check and re-init some states
30
            self.self_check_during_idle()
31

32
        self.process_disagg_prefill_inflight_queue()
33

34
        # Run sample of the current batch
35
        # It depends on the result of the last batch (e.g., grammar), so we run it after the last batch is processed.
36
        self.launch_batch_sample_if_needed(batch_result)
37

38
        # Update last_batch
39
        self.last_batch = batch

• https://github.com/zhaochenyang20/Awesome-ML-SYS-Tutorial/blob/main/sglang/code-walk-through/multimodal_request_lifecycle.md
  • For multi-model and RoPE handling, see MRotaryEmbedding and its get_rope_index;

Chunked Prefill

• https://zhuanlan.zhihu.com/p/718715866
  • Decode costs more per token than prefill. Prefill can batch its work, whereas decode proceeds token by token and depends heavily on KV-cache I/O.
  • Split prefill prompts into chunks and insert decode work during staging/I/O, placing its single token in unused capacity within prefill chunks. This reduces prefill efficiency through extra stages and KV I/O, but provides more resources and time for decoding;

ML Systems Fundamentals

• torch-memory-savor
  • CUDA graphs: manually compile several Torch operations into a CUDA graph to reduce GPU<->CPU instruction transfers, avoiding CPU bottlenecks in complex topologies such as FSDP. This resembles arranging fixed operators into DAGs in Spark/Ray. It is somewhat like torch.compile, but torch.compile targets developers and general acceleration rather than CUDA alone.
  • torch-memory-savor manages SGLang’s CUDA memory to maintain stable GPU memory for CUDA graphs. It uses low-level CUDA APIs to map and manage logical and physical memory, with pause/resume/malloc methods. I do not fully understand why physical memory must be released so frequently. Would allocating one large region, reusing it, and garbage-collecting not be better?
• nccl(nvidia collective communication library)
  • AllReduce: reduce data from all GPUs, for example by summing, then broadcast to all GPUs
  • Broadcast: send data from one source GPU to every other GPU
  • Reduce: reduce all GPUs’ data onto one destination GPU
  • AllGather: collect data from all GPUs and distribute it to every GPU, dim * world_size
  • ReduceScatter: reduce data, then scatter the result across GPUs
  • Send/Recv: point-to-point communication
  • AllToAll: distribute data to all GPUs
  • send/recv
  • isend/irecv/batch_isend_irecv: asynchronous; wait must be called
• Communication algorithms: ring versus tree
  • ring: all-reduce requires 2*(n-1) communication steps, n-1 reductions plus n-1 broadcasts
  • tree: all-reduce requires 2*log(N) communication steps
  • Since NCCL 2.4, cross-node communication with many nodes uses Double Binary Tree. Compared with a traditional tree algorithm, two complementary binary trees balance communication overhead.
• rl-memory-management(https://hebiao064.github.io/rl-memory-management)
  • Not studied yet; will return after learning some RL