Refactor Mamba chunked prefill

megatron/core/inference/batch_dimensions_utils.py

@@ -25,7 +25,6 @@ class InferenceBatchDimensions:
         token_count : number of total input tokens
         prefill_req_count : number of prefill requests
         decode_req_count : number of decode requests
-        has_explicit_chunked_prefill_req : whether the batch has an explicit chunked prefill request
 
     The batch dimensions are ordered by token_count, then by prefill_req_count,
     then by decode_req_count.
@@ -35,7 +34,6 @@ class InferenceBatchDimensions:
     token_count: int = 0
     prefill_req_count: int = 0
     decode_req_count: int = 0
-    has_explicit_chunked_prefill_req: bool = False
 
     def __str__(self):
         """
@@ -55,9 +53,6 @@ def is_applicable_for_batch_dim(
         for prefill or decode requests. Otherwise, prefill slots
         can only be used for prefill requests.
         """
-        if real_batch_dim.has_explicit_chunked_prefill_req != self.has_explicit_chunked_prefill_req:
-            return False
-
         if real_batch_dim.prefill_req_count == 0:
             return (
                 self.token_count >= real_batch_dim.token_count
@@ -104,42 +99,25 @@ def is_valid(self, max_requests: int, max_sequence_length: int) -> bool:
         if self.token_count > self.prefill_req_count * max_sequence_length + self.decode_req_count:
             return False
 
-        # Check if there is an invalid chunked prefill request.
-        if self.prefill_req_count == 0 and self.has_explicit_chunked_prefill_req:
-            return False
-
         return True
 
     def __hash__(self):
         """
         Returns a hash of the batch dimension.
         In cuda graph quick matching, the batch dimension is used as a key in a dictionary.
         """
-        return hash(
-            (
-                self.token_count,
-                self.prefill_req_count,
-                self.decode_req_count,
-                self.has_explicit_chunked_prefill_req,
-            )
-        )
+        return hash((self.token_count, self.prefill_req_count, self.decode_req_count))
 
     def __eq__(self, other: "InferenceBatchDimensions") -> bool:
         """
         Checks if this batch dimension is equal to another batch dimension.
         """
         if other is None:
             return False
-        return (
-            self.token_count,
-            self.prefill_req_count,
-            self.decode_req_count,
-            self.has_explicit_chunked_prefill_req,
-        ) == (
+        return (self.token_count, self.prefill_req_count, self.decode_req_count) == (
             other.token_count,
             other.prefill_req_count,
             other.decode_req_count,
-            other.has_explicit_chunked_prefill_req,
         )
 
     @property
@@ -154,6 +132,7 @@ def adjust_batch_dims_for_expert_parallelism(
         local_batch_dims,
         strict: bool,
         decode_only_cuda_graphs: bool,
+        explicit_chunked_prefill: bool,
         ep_group: Optional[torch.distributed.ProcessGroup] = None,
     ) -> Optional["InferenceBatchDimensions"]:
         """Adjusted cuda graph batch dimensions for expert parallelism.
@@ -163,6 +142,7 @@ def adjust_batch_dims_for_expert_parallelism(
             local_batch_dims: The local batch dimensions to adjust.
             strict: Whether to use strict matching for batch dimensions.
             decode_only_cuda_graphs: Whether CUDA graphs are only used for decode steps.
+            explicit_chunked_prefill: Whether chunked prefill is enabled with explicit requests
             ep_group: Optional expert parallel process group. If None, uses global parallel state.
                       When using different EP sizes for inference vs training, pass the
                       inference EP group explicitly.
@@ -176,13 +156,11 @@ def adjust_batch_dims_for_expert_parallelism(
             return local_batch_dims
         # all reduce local work across expert model parallel group
 
-        has_explicit_chunked_prefill_req = local_batch_dims.has_explicit_chunked_prefill_req
         is_non_decode = local_batch_dims.prefill_req_count > 0
         sync_tensor = torch.tensor(
             [
                 local_batch_dims.token_count,
                 int(is_non_decode),
-                int(has_explicit_chunked_prefill_req),
                 local_batch_dims.prefill_req_count,
                 local_batch_dims.decode_req_count,
             ],
@@ -194,7 +172,6 @@ def adjust_batch_dims_for_expert_parallelism(
 
         sync_tensor = sync_tensor.cpu()
         is_any_ep_rank_in_non_decode = sync_tensor[1].item() == 1
-        any_ep_rank_has_explicit_chunked_prefill_req = sync_tensor[2].item() == 1
 
         # We force eager mode for scenarios where some ranks will run with CUDA graphs
         # while others will not. Without this check, the all-to-all communication in the
@@ -204,28 +181,23 @@ def adjust_batch_dims_for_expert_parallelism(
         #   1. If we only allow decode CUDA graphs but some ranks are running non-decode batches
         #   2. Some ranks are running explicit chunked prefill requests
         #       (graphs are not recorded for batches with explicit chunked prefill requests)
-        if (
-            decode_only_cuda_graphs and is_any_ep_rank_in_non_decode
-        ) or any_ep_rank_has_explicit_chunked_prefill_req:
+        if is_any_ep_rank_in_non_decode and (decode_only_cuda_graphs or explicit_chunked_prefill):
             return None  # indicate no match, run in eager mode
 
-        assert not has_explicit_chunked_prefill_req
-
         # If strict matching is enabled, we sync the request counts across EP ranks
         # to ensure the graph captures the maximum needed capacity.
         # TODO(ksanthanam): Add functional test for this scenario
         adjusted_prefill_req_count = (
-            int(sync_tensor[3].item()) if strict else local_batch_dims.prefill_req_count
+            int(sync_tensor[2].item()) if strict else local_batch_dims.prefill_req_count
         )
         adjusted_decode_req_count = (
-            int(sync_tensor[4].item()) if strict else local_batch_dims.decode_req_count
+            int(sync_tensor[3].item()) if strict else local_batch_dims.decode_req_count
         )
 
         adjusted_batch_dim = InferenceBatchDimensions(
             token_count=int(sync_tensor[0].item()),
             prefill_req_count=adjusted_prefill_req_count,
             decode_req_count=adjusted_decode_req_count,
-            has_explicit_chunked_prefill_req=False,
         )
         return adjusted_batch_dim
 
@@ -462,6 +434,7 @@ def match_graph_config(
         cuda_graph_batch_dimensions_list: List[InferenceBatchDimensions],
         strict: bool = False,
         decode_only_cuda_graphs: bool = False,
+        explicit_chunked_prefill: bool = False,
         ep_group: Optional[torch.distributed.ProcessGroup] = None,
     ) -> Optional[InferenceBatchDimensions]:
         """
@@ -475,6 +448,7 @@ def match_graph_config(
             decode_only_cuda_graphs: Used by expert parallel matching. If this is true,
             and one of the EP ranks is running a non-decode step, we elect to run in
             eager mode instead of matching a decode-only cuda graph.
+            explicit_chunked_prefill: Whether chunked prefill is enabled with explicit requests
             ep_group: Optional expert parallel process group. If None, uses global parallel state.
                       When using different EP sizes for inference vs training, pass the
                       inference EP group explicitly.
@@ -490,6 +464,7 @@ def match_graph_config(
             real_batch_dim,
             strict=strict,
             decode_only_cuda_graphs=decode_only_cuda_graphs,
+            explicit_chunked_prefill=explicit_chunked_prefill,
             ep_group=ep_group,
         )
 
@@ -499,6 +474,9 @@ def match_graph_config(
             # in that case, all ranks have to run in eager mode
             return None
 
+        if explicit_chunked_prefill and real_batch_dim.prefill_req_count > 0:
+            return None
+
         # first filter out batch dimensions with smaller token count, prefill req count,
         # or decode req count, as they are not applicable
         graph_batch_dims_applicable = [

megatron/core/inference/contexts/attention_context/mamba_metadata.py

@@ -57,8 +57,8 @@ def __init__(self, max_requests: int, max_tokens: int):
         )
 
         # Tuple of (
-        #   total prefill sequence length excluding chunked prefill,
-        #   chunked prefill sequence length
+        #   chunked prefill sequence length,
+        #   total regular prefill sequence length
         # )
         self._device_chunked_prefill_buffer = torch.zeros(
             (2,), dtype=torch.int32, device=self.device
@@ -101,6 +101,7 @@ def update(
         cu_seqlens: torch.Tensor,
         batch_dimensions: InferenceBatchDimensions,
         padded_batch_dimensions: InferenceBatchDimensions,
+        enable_chunked_prefill: bool,
     ) -> None:
         """
         Updates the dedicated CUDA graph mapping tensor with the indices
@@ -109,20 +110,20 @@ def update(
         Args:
             active_mamba_indices (Tensor): Tensor containing the Mamba slot indices
                                            for active requests.
-            num_active_requests (int): The number of active requests.
+            token_to_request_idx (Tensor): Map from token index to request index.
+            cu_seqlens (Tensor): Cumulative sequence lengths.
+            batch_dimensions (InferenceBatchDimensions): Dimensions of the current batch.
+            padded_batch_dimensions (InferenceBatchDimensions): Dimensions of the padded batch.
+                enable_chunked_prefill (bool): Whether chunked prefill is enabled
         """
         real_decode_count = batch_dimensions.decode_req_count
         real_prefill_count = batch_dimensions.prefill_req_count
-        real_token_count = batch_dimensions.token_count
-        has_explicit_chunked_prefill_req = batch_dimensions.has_explicit_chunked_prefill_req
 
         padded_decode_count = padded_batch_dimensions.decode_req_count
         padded_prefill_count = padded_batch_dimensions.prefill_req_count
         padded_token_count = padded_batch_dimensions.token_count
-        assert (
-            has_explicit_chunked_prefill_req
-            == padded_batch_dimensions.has_explicit_chunked_prefill_req
-        )
+
+        has_chunked_prefill_req = enable_chunked_prefill and real_prefill_count > 0
 
         if padded_decode_count > 0:
             # Update decode indices
@@ -135,10 +136,12 @@ def update(
 
         # Determine if we have a chunked prefill request and adjust counts for regular prefill
         regular_prefill_count = real_prefill_count
-        if has_explicit_chunked_prefill_req:
-            # The last prefill request is the chunked one
+        chunked_req_idx = -1
+
+        if has_chunked_prefill_req:
+            # The first prefill request is the chunked one
             regular_prefill_count -= 1
-            chunked_req_idx = real_decode_count + regular_prefill_count
+            chunked_req_idx = real_decode_count
 
             # Update chunked prefill indices
             self._batch_indices_chunked_prefill_buffer[0] = active_mamba_indices[chunked_req_idx]
@@ -149,41 +152,61 @@ def update(
         if padded_prefill_count > 0:
             # Update prefill indices (excluding chunked prefill from regular prefill buffer)
             if regular_prefill_count > 0:
+                # If chunked prefill exists, regular prefills start after it.
+                # If no chunked prefill, regular prefills start at real_decode_count.
+                start_idx = real_decode_count + (1 if has_chunked_prefill_req else 0)
+
                 self._batch_indices_prefill_buffer[:regular_prefill_count].copy_(
-                    active_mamba_indices[
-                        real_decode_count : real_decode_count + regular_prefill_count
-                    ]
+                    active_mamba_indices[start_idx : start_idx + regular_prefill_count]
                 )
 
             if padded_prefill_count > regular_prefill_count:
                 self._batch_indices_prefill_buffer[regular_prefill_count:padded_prefill_count] = -1
 
             self.batch_indices_prefill = self._batch_indices_prefill_buffer[:padded_prefill_count]
 
-            # Update seq_idx
-            end_regular_prefill_token_idx = cu_seqlens[real_decode_count + regular_prefill_count]
+            # Update seq_idx for regular prefills
+            # If chunked prefill exists, we need to skip its tokens in seq_idx
+
+            # Index where regular prefills end in the batch (decode + chunked + regular)
+            end_regular_prefill_req_idx = (
+                real_decode_count + regular_prefill_count + (1 if has_chunked_prefill_req else 0)
+            )
+            end_regular_prefill_token_idx = cu_seqlens[end_regular_prefill_req_idx]
+
+            # Index where regular prefills start
+            start_regular_prefill_req_idx = real_decode_count + (
+                1 if has_chunked_prefill_req else 0
+            )
+            start_regular_prefill_token_idx = cu_seqlens[start_regular_prefill_req_idx]
 
-            # The length of tokens belonging to regular prefill requests (excluding decode tokens)
-            seq_len = end_regular_prefill_token_idx - real_decode_count
+            # The length of tokens belonging to regular prefill requests
+            seq_len = end_regular_prefill_token_idx - start_regular_prefill_token_idx
 
             if seq_len > 0:
+                # We subtract start_regular_prefill_req_idx to normalize request IDs to
+                # 0-based relative to this buffer
                 self._seq_idx_buffer[:, :seq_len].copy_(
-                    token_to_request_idx[real_decode_count:end_regular_prefill_token_idx]
-                    - real_decode_count
+                    token_to_request_idx[
+                        start_regular_prefill_token_idx:end_regular_prefill_token_idx
+                    ]
+                    - start_regular_prefill_req_idx
                 )
 
             if padded_token_count > seq_len:
                 self._seq_idx_buffer[:, seq_len:padded_token_count] = -1
             self.seq_idx = self._seq_idx_buffer[:, :padded_token_count]
 
-            # Update cu_seqlens
+            # Update cu_seqlens for regular prefill requests
             self._cu_seqlens_buffer[0] = 0
             if regular_prefill_count > 0:
+                # Copy cu_seqlens for regular prefill requests and normalize by
+                # subtracting the start token index
+                start_req_idx = real_decode_count + (1 if has_chunked_prefill_req else 0)
+                end_req_idx = start_req_idx + regular_prefill_count
+
                 self._cu_seqlens_buffer[1 : regular_prefill_count + 1].copy_(
-                    cu_seqlens[
-                        real_decode_count + 1 : real_decode_count + regular_prefill_count + 1
-                    ]
-                    - real_decode_count
+                    cu_seqlens[start_req_idx + 1 : end_req_idx + 1] - cu_seqlens[start_req_idx]
                 )
 
             # Pad the rest with the last value (effectively length 0 segments)
@@ -195,19 +218,32 @@ def update(
 
         if padded_decode_count > 0 and padded_prefill_count > 0:
             self._device_decode_prefill_buffer[0] = real_decode_count
-            self._device_decode_prefill_buffer[1] = regular_prefill_count
+            # This describes the number of items in the prefill tensor relative to the
+            # decode tensor. If chunked prefill is present, it is included in the
+            # "prefill" part of the main split.
+            self._device_decode_prefill_buffer[1] = regular_prefill_count + (
+                1 if has_chunked_prefill_req else 0
+            )
             self.device_decode_prefill = self._device_decode_prefill_buffer
 
-        # If using chunked prefill for this batch, store the number of regular prefill tokens
-        # and the number of tokens in the chunked prefill request
-        if has_explicit_chunked_prefill_req:
+        # If using chunked prefill for this batch, store the number of chunked tokens
+        # and the number of regular prefill tokens
+        if has_chunked_prefill_req:
+            # Chunked request is the first prefill request (index real_decode_count)
             chunked_prefill_token_count = (
-                cu_seqlens[real_decode_count + real_prefill_count]
-                - cu_seqlens[real_decode_count + real_prefill_count - 1]
+                cu_seqlens[real_decode_count + 1] - cu_seqlens[real_decode_count]
             )
-            assert self.cu_seqlens is not None
-            self._device_chunked_prefill_buffer[0] = self.cu_seqlens[regular_prefill_count]
-            self._device_chunked_prefill_buffer[1] = chunked_prefill_token_count
+
+            # Regular prefill tokens are everything after the chunked request tokens
+            regular_prefill_token_count = 0
+            if regular_prefill_count > 0:
+                regular_prefill_token_count = (
+                    cu_seqlens[real_decode_count + 1 + regular_prefill_count]
+                    - cu_seqlens[real_decode_count + 1]
+                )
+
+            self._device_chunked_prefill_buffer[0] = chunked_prefill_token_count
+            self._device_chunked_prefill_buffer[1] = regular_prefill_token_count
             self.device_chunked_prefill = self._device_chunked_prefill_buffer
 
     def allocate_slot(self) -> Optional[int]: