Cuda graph padding (#174)

* Support padding for cuda graph

* Fix select graph size and default max cuda graph size

* Set default to 512

* Fix issues

* Support moonlight

Author: gty111

README.md

@@ -153,7 +153,7 @@ python benchmarks/evaluate_MMLU_pro.py --model $MODEL
 
 ## Supported Models
 
-- Kimi Series: K2-Base, K2-Instruct
+- Kimi Series: Moonlight, K2-Base, K2-Instruct
 - DeepSeek Series: DeepSeek R1, DeepSeek V3, DeepSeek V2
 - Qwen Series: Qwen3 VL, Qwen3, Qwen2.5 VL, Qwen2.5, Qwen2
 - Llama Series: Llama3.2, Llama3.1, Llama3, Llama2 and deepseek-coder

gllm/entrypoints/api_server.py

@@ -190,8 +190,13 @@ async def run_server(args):
     parser.add_argument(
         "--max-cuda-graph-bs",
         type=int,
-        help="Maximum batch size for cuda graph",
-        default=32,
+        help=(
+            "Maximum batch size for CUDA graph capture. "
+            "Larger values allow more decode batches to benefit from CUDA graphs "
+            "but increase startup time and GPU memory usage during graph capture. "
+            "Default: 512."
+        ),
+        default=512,
     )
     # Parallelism
     parser.add_argument("--pp", type=int, help="Number of pipeline stages", default=1)

gllm/input_data.py

@@ -316,6 +316,56 @@ def _cal_mla_metadata(self, seqs: List[Sequence]):
                     dtype=torch.int32,
                 )
 
+    def pad_for_cuda_graph(self, padded_size: int):
+        """Pad input buffers to padded_size using dummy values.
+
+        This enables CUDA graph replay for a fixed batch size (a power-of-two
+        bucket) even when the actual number of decode tokens is smaller.
+
+        The dummy tokens write their KV entries to memory_manager.dummy_page,
+        which is permanently reserved and never used by real sequences.
+
+        Returns:
+            num_real_tokens (int): the actual (unpadded) token count, so that
+            the caller can slice output_hidden_states[:num_real_tokens] when
+            computing logits after graph replay.
+        """
+        assert self.use_buffer, "pad_for_cuda_graph requires use_buffer=True"
+        num_real_tokens = self.tokens_cpu.shape[0]
+        if num_real_tokens >= padded_size:
+            return num_real_tokens
+
+        dummy_page = self.memory_manager.dummy_page
+        dummy_slot = dummy_page * self.page_size  # slot index within dummy page
+
+        num_pad = padded_size - num_real_tokens
+
+        # tokens: pad with 0
+        self.tokens[num_real_tokens:padded_size].zero_()
+        # positions: pad with 0
+        self.positions[num_real_tokens:padded_size].zero_()
+        # mrope_positions: pad with 0
+        self.mrope_positions[:, num_real_tokens:padded_size].zero_()
+        # slot_mapping: pad with dummy slot so writes go to the reserved page
+        self.slot_mapping[num_real_tokens:padded_size].fill_(dummy_slot)
+        # seq_lens: pad with 1 (avoid division-by-zero in attention kernels)
+        self.seq_lens[len(self.seqs):len(self.seqs) + num_pad].fill_(1)
+        # block_table: pad rows with dummy_page
+        self.block_table[len(self.seqs):len(self.seqs) + num_pad].fill_(dummy_page)
+        # query_start_loc: continue the cumulative sum — each dummy token counts
+        # as 1 query token, so the padded entries are last_loc+1, last_loc+2, ...
+        last_loc = self.query_start_loc[len(self.seqs)]
+        self.query_start_loc[len(self.seqs) + 1:len(self.seqs) + num_pad + 1].copy_(
+            last_loc + torch.arange(1, num_pad + 1, dtype=torch.int32)
+        )
+
+        if self.use_mla:
+            # Pad decode_seq_lens for the dummy sequences so that MLA decode
+            # kernels see a valid (non-zero) sequence length for every row.
+            self.decode_seq_lens[len(self.seqs):len(self.seqs) + num_pad].fill_(1)
+
+        return num_real_tokens
+
     def _set_mla_metadata(self):
         if self.num_prefills > 0:
             self.prefill_query_start_loc[

gllm/memory_manager.py

@@ -86,7 +86,7 @@ def __init__(
         self.vocab_size = vocab_size
         self.use_mla = use_mla
 
-    def init(self, segment_cls=Segment):
+    def init(self, segment_cls=Segment, reserve_dummy_page: bool = False):
         free_mem_size, _ = torch.cuda.mem_get_info()
         num_max_pages = free_mem_size // self.get_sizeof_KV_per_page()
         num_pages = int(num_max_pages * self.gpu_memory_util)
@@ -113,6 +113,12 @@ def init(self, segment_cls=Segment):
             self.use_mla,
         )
 
+        # Reserve a dedicated dummy page for CUDA graph padding only when
+        # CUDA graphs are enabled.  This page is never returned to normal use,
+        # so real sequences will never overwrite it, and padding dummy tokens
+        # can safely write here.
+        self.dummy_page: int = self.segment.allocate() if reserve_dummy_page else None
+
         self.kv_cache_dtype = "auto"
         self.k_scale = torch.tensor(1.0, dtype=torch.float32)
         self.v_scale = self.k_scale
@@ -182,8 +188,8 @@ class PrefixMemoryManager(MemoryManager):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
-    def init(self):
-        super().init(segment_cls=PrefixSegment)
+    def init(self, reserve_dummy_page: bool = False):
+        super().init(segment_cls=PrefixSegment, reserve_dummy_page=reserve_dummy_page)
 
         # for prefill stage
         self.num_allocated_pages = 0

gllm/model_runner.py

@@ -127,7 +127,9 @@ def __init__(
         self.disable_cuda_graph = disable_cuda_graph
         self.max_cuda_graph_bs = max_cuda_graph_bs
         self.size_to_graph: Dict[int, torch.cuda.CUDAGraph] = dict()
-        self.capture_sizes = list(range(self.max_cuda_graph_bs, 0, -1))
+        # Use power-of-two bucket sizes to reduce the number of captured graphs.
+        # At runtime the actual batch is padded up to the nearest bucket.
+        self.capture_sizes = self._build_capture_sizes(self.max_cuda_graph_bs)
 
         # max length
         self.model_max_length = self.resolve_model_max_length(model_max_length)
@@ -142,6 +144,26 @@ def resolve_model_max_length(self, model_max_length):
         logger.info(f"Model max length: {model_max_length}")
         return model_max_length
 
+    @staticmethod
+    def _build_capture_sizes(max_bs: int):
+        """Return power-of-two bucket sizes up to max_bs, in descending order.
+
+        For example, max_bs=20 → [20, 16, 8, 4, 2, 1].
+        We always include 1 as a floor bucket.
+        """
+        if max_bs <= 0:
+            return []
+        sizes = []
+        s = 1
+        while s <= max_bs:
+            sizes.append(s)
+            s *= 2
+        # If max_bs is not itself a power of two, add it as the top bucket so
+        # that batches of exactly max_bs can still use CUDA graph.
+        if sizes[-1] != max_bs:
+            sizes.append(max_bs)
+        return list(reversed(sizes))
+
     def init(self, mp_load_progress=None):
         self.model = self.model_loader.load_model(mp_load_progress)
         memory_manager_cls = (
@@ -171,8 +193,9 @@ def init(self, mp_load_progress=None):
         self.output_residual = torch.zeros((self.max_num_batched_tokens, self.hidden_size))
         # Profile run
         self.profile_run()
-        # Init KV cache at last
-        self.memory_manager.init()
+        # Init KV cache at last; only reserve the dummy page when CUDA graphs
+        # are actually enabled so we don't waste memory otherwise.
+        self.memory_manager.init(reserve_dummy_page=not self.disable_cuda_graph)
 
         if not self.disable_cuda_graph:
             self.capture_graph()
@@ -373,7 +396,7 @@ def profile_run(self):
     def capture_graph(self):
         iterator = self.capture_sizes
         if get_local_rank() == 0:
-            # logger.info(f"Capturing cuda graph for sizes {self.capture_sizes}")
+            logger.info(f"Capturing CUDA graphs for bucket sizes: {list(reversed(self.capture_sizes))}")
             iterator = tqdm(self.capture_sizes, desc="Capturing CUDA Graphs", ncols=100)
         memory_pool = torch.cuda.graph_pool_handle()
         for size in iterator:
@@ -426,9 +449,22 @@ def check_decode_batch(self):
     @torch.inference_mode()
     def step_once(self):
         num_cal_tokens = self.input_data.tokens_cpu.shape[0]
-        # Only decode batch use cuda graph
-        if self.check_decode_batch() and num_cal_tokens in self.size_to_graph:
-            self.size_to_graph[num_cal_tokens].replay()
+        # Only pure decode batches use CUDA graph.
+        if self.check_decode_batch():
+            # Find the smallest captured bucket >= actual batch size.
+            padded_size = None
+            for bucket in self.capture_sizes:
+                if bucket >= num_cal_tokens:
+                    padded_size = bucket
+            if padded_size is not None and padded_size in self.size_to_graph:
+                # Pad input buffers to the bucket size with dummy values, then
+                # replay the pre-captured graph.
+                num_real_tokens = self.input_data.pad_for_cuda_graph(padded_size)
+                self.size_to_graph[padded_size].replay()
+                # After replay, use only the real-token slice for logits.
+                num_cal_tokens = num_real_tokens
+            else:
+                self.forward()
         else:
             self.forward()
         if is_last_pp_rank():

gllm/models/deepseek_v2.py

@@ -157,7 +157,18 @@ def __init__(self, layer_id: int, config):
         self.kv_lora_rank = config.kv_lora_rank
         self.rope_theta = getattr(config, "rope_theta", 10000)
         self.max_poistion_embeddings = getattr(config, "max_position_embeddings", 8192)
-        self.rope_scaling = getattr(config, "rope_scaling", None)
+        rope_scaling = getattr(config, "rope_scaling", None)
+        if rope_scaling is None:
+            self.rope_scaling = {
+                "factor": 1.0,
+                "original_max_position_embeddings": self.max_poistion_embeddings,
+            }
+        else:
+            self.rope_scaling = dict(rope_scaling)
+            self.rope_scaling.setdefault("factor", 1.0)
+            self.rope_scaling.setdefault(
+                "original_max_position_embeddings", self.max_poistion_embeddings
+            )
 
         if self.q_lora_rank is not None:
             self.q_a_proj = ReplicatedLinear(
@@ -286,7 +297,18 @@ def __init__(self, layer_id: int, config):
         self.kv_lora_rank = config.kv_lora_rank
         self.rope_theta = getattr(config, "rope_theta", 10000)
         self.max_poistion_embeddings = getattr(config, "max_position_embeddings", 8192)
-        self.rope_scaling = getattr(config, "rope_scaling", None)
+        rope_scaling = getattr(config, "rope_scaling", None)
+        if rope_scaling is None:
+            self.rope_scaling = {
+                "factor": 1.0,
+                "original_max_position_embeddings": self.max_poistion_embeddings,
+            }
+        else:
+            self.rope_scaling = dict(rope_scaling)
+            self.rope_scaling.setdefault("factor", 1.0)
+            self.rope_scaling.setdefault(
+                "original_max_position_embeddings", self.max_poistion_embeddings
+            )
         self.layer_id = layer_id
 
         if self.q_lora_rank is not None: