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 = [