Move inference context bookkeeping to CPU with ContextGPUView

docs/user-guide/features/megatron_fsdp.md

@@ -18,7 +18,7 @@
 
 ### 🧩 Compatibility
 
-- PyTorch **[DeviceMesh](https://docs.pytorch.org/docs/stable/distributed.html#devicemesh)**, **[DTensor](https://docs.pytorch.org/docs/stable/distributed.tensor.html)**, and **[Distributed Checkpoint (DCP)](https://docs.pytorch.org/docs/stable/distributed.checkpoint.html)**
+- PyTorch **[DeviceMesh](https://docs.pytorch.org/docs/2.11/distributed.html#torch.distributed.device_mesh.DeviceMesh)**, **[DTensor](https://docs.pytorch.org/docs/stable/distributed.tensor.html)**, and **[Distributed Checkpoint (DCP)](https://docs.pytorch.org/docs/stable/distributed.checkpoint.html)**
 - **[Megatron Core](https://github.com/NVIDIA/Megatron-LM)**
 - **[TransformerEngine](https://github.com/NVIDIA/TransformerEngine)**
 - **[NVIDIA NeMo Framework Container](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/nemo)**
@@ -605,4 +605,4 @@ NCCL (`v2.27+`) supports symmetric allocation or registration for communicators
 - **Copy-Engine (CE) Collectives**: Instead of using SMs (or CTAs) for common non-computational collectives like AG in Megatron-FSDP, copy engines are instead used to perform all-gather collectives, dedicating SM resources to compute and reduction during FSDP. Requires NCCL `v2.28+`.
 - **High-Precision Reduction**: When training large models, high-precision gradient reduction and accumulation is desired for accuracy and convergence, but communicating FP32 gradients is expensive. With symmetric registration, FP32 accumulators enable gradients to be reduced in FP32 but communicated in BF16, which decreases gradient RS communication latency while maintaining high accuracy during training. Megatron-FSDP supports FP32 main gradient accumulation but BF16 gradient communication, customizable through `megatron_fsdp.MixedPrecisionPolicy`.
 
-These optimizations significantly reduce SM resource contention for overlapped compute and communication kernels in FSDP. Symmetric registration, allocation, and pooling is also supported in PyTorch: [`torch.distributed._symmetric_memory`](https://docs.pytorch.org/docs/stable/symmetric_memory.html).
+These optimizations significantly reduce SM resource contention for overlapped compute and communication kernels in FSDP. Symmetric registration, allocation, and pooling is also supported in PyTorch: [`torch.distributed._symmetric_memory`](https://docs.pytorch.org/docs/stable/symmetric_memory.html).

megatron/core/inference/batch_dimensions_utils.py

@@ -141,7 +141,13 @@ def req_count(self) -> int:
 
     @staticmethod
     def adjust_batch_dims_for_expert_parallelism(
-        local_batch_dims, ep_group: Optional[torch.distributed.ProcessGroup] = None
+        local_batch_dims,
+        strict: bool = False,
+        decode_only_cuda_graphs: bool = True,
+        smallest_non_decode_cuda_graph_size: int = 0,
+        ep_group: Optional[torch.distributed.ProcessGroup] = None,
+        num_speculative_tokens: int = 0,
+        ep_zmq_communicator=None,
     ) -> Optional["InferenceBatchDimensions"]:
         """Adjust CUDA graph batch dimensions for expert parallelism.
 
@@ -152,7 +158,16 @@ def adjust_batch_dims_for_expert_parallelism(
 
         Args:
             local_batch_dims: The local batch dimensions to adjust.
-            ep_group: Expert parallel process group.
+            strict: Whether to use strict matching for batch dimensions.
+            decode_only_cuda_graphs: Whether CUDA graphs are only used for decode steps.
+            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.
+            ep_zmq_communicator: Optional AsyncZMQCommunicator over the EP group. When
+                      provided, the cross-rank MAX reduction runs on the CPU via ZMQ
+                      (no GPU kernel, no H2D/D2H), avoiding a per-step NCCL AllReduce
+                      on the compute stream. When absent, falls back to
+                      torch.distributed.all_reduce on a GPU tensor.
 
         Returns:
             InferenceBatchDimensions with max token count, or None for eager mode.
@@ -162,23 +177,77 @@ def adjust_batch_dims_for_expert_parallelism(
             return local_batch_dims
 
         is_non_decode = local_batch_dims.prefill_req_count > 0
-        sync_tensor = torch.tensor(
-            [local_batch_dims.token_count, int(is_non_decode)],
-            dtype=torch.int32,
-            device=torch.cuda.current_device(),
-        )
-        torch.distributed.all_reduce(sync_tensor, op=torch.distributed.ReduceOp.MAX, group=ep_group)
-        sync_tensor = sync_tensor.cpu()
 
-        if sync_tensor[1].item() == 1:
+        if ep_zmq_communicator is not None:
+            # CPU-only sync via ZMQ: avoids a NCCL AllReduce kernel on the
+            # compute stream plus the H2D/D2H pair that sandwiches it.
+            (max_token_count, max_is_non_decode, max_prefill_count, max_decode_count) = (
+                ep_zmq_communicator.sync_all_reduce_max(
+                    local_batch_dims.token_count,
+                    int(is_non_decode),
+                    local_batch_dims.prefill_req_count,
+                    local_batch_dims.decode_req_count,
+                )
+            )
+        else:
+            sync_tensor = torch.tensor(
+                [
+                    local_batch_dims.token_count,
+                    int(is_non_decode),
+                    local_batch_dims.prefill_req_count,
+                    local_batch_dims.decode_req_count,
+                ],
+                dtype=torch.int32,
+                device=torch.cuda.current_device(),
+            )
+            torch.distributed.all_reduce(
+                sync_tensor, op=torch.distributed.ReduceOp.MAX, group=ep_group
+            )
+            sync_tensor = sync_tensor.cpu()
+            max_token_count = int(sync_tensor[0].item())
+            max_is_non_decode = int(sync_tensor[1].item())
+            max_prefill_count = int(sync_tensor[2].item())
+            max_decode_count = int(sync_tensor[3].item())
+
+        is_any_ep_rank_in_non_decode = max_is_non_decode == 1
+
+        if is_any_ep_rank_in_non_decode and decode_only_cuda_graphs:
             return None  # any rank has prefill → eager mode
 
-        return InferenceBatchDimensions(
-            token_count=int(sync_tensor[0].item()),
-            prefill_req_count=local_batch_dims.prefill_req_count,
-            decode_req_count=local_batch_dims.decode_req_count,
+        adjusted_token_count = max_token_count
+
+        # Sync request counts across EP ranks when strict matching is enabled
+        # or when speculative tokens are used.  With speculative tokens,
+        # decode-only graphs have token counts of decode_req_count * (spec+1)
+        # which creates a different granularity than mixed graphs (raw sizes).
+        # Without syncing, decode-only ranks and prefill ranks search different
+        # graph pools and may pick graphs with different token counts.
+        sync_request_counts = strict or (
+            is_any_ep_rank_in_non_decode and num_speculative_tokens > 0
+        )
+        adjusted_prefill_req_count = (
+            max_prefill_count if sync_request_counts else local_batch_dims.prefill_req_count
+        )
+        adjusted_decode_req_count = (
+            max_decode_count if sync_request_counts else local_batch_dims.decode_req_count
         )
 
+        # When any EP rank has prefill requests (non-strict mode), elevate
+        # the token count to be >= the smallest prefill/mixed cuda graph.
+        # This ensures decode-only ranks don't match a fine-grained decode
+        # graph while prefill ranks match a coarser mixed graph, which would
+        # produce inconsistent token counts across EP ranks.
+        if is_any_ep_rank_in_non_decode and not strict:
+            adjusted_token_count = max(adjusted_token_count, smallest_non_decode_cuda_graph_size)
+
+        adjusted_batch_dim = InferenceBatchDimensions(
+            token_count=adjusted_token_count,
+            prefill_req_count=adjusted_prefill_req_count,
+            decode_req_count=adjusted_decode_req_count,
+        )
+
+        return adjusted_batch_dim
+
 
 class CUDAGraphBatchDimensionBuilder:
     """Builder for creating and managing CUDA graph batch dimensions.
@@ -460,8 +529,12 @@ def add_if_valid(token_count: int, prefill_req_count: int, decode_req_count: int
     def match_graph_config(
         real_batch_dim: InferenceBatchDimensions,
         cuda_graph_batch_dimensions_list: List[InferenceBatchDimensions],
+        smallest_non_decode_cuda_graph_size: int = 0,
         strict: bool = False,
+        decode_only_cuda_graphs: bool = True,
         ep_group: Optional[torch.distributed.ProcessGroup] = None,
+        num_speculative_tokens: int = 0,
+        ep_zmq_communicator=None,
         match_ep_token_counts: bool = True,
     ) -> Optional[InferenceBatchDimensions]:
         """
@@ -472,9 +545,16 @@ def match_graph_config(
             cuda_graph_batch_dimensions_list: List of available CUDA graph batch dimensions
             strict: If False, prefill slots can be used for prefill or decode requests.
                    If True, prefill slots can only be used for prefill requests.
+            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.
             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.
+            ep_zmq_communicator: Optional AsyncZMQCommunicator over the EP group. When
+                      provided, batch-dimension MAX reduction uses a CPU-only ZMQ sync
+                      instead of a GPU NCCL AllReduce. Forwarded to
+                      adjust_batch_dims_for_expert_parallelism.
             match_ep_token_counts: If True (default), token counts are synced across EP ranks via
                 all-reduce-max so all ranks select the same CUDA graph. Set to False when the
                 dispatcher handles per-rank token variation internally (e.g. AGV/RSV in the NVLS
@@ -491,9 +571,19 @@ def match_graph_config(
             # NCCL dispatcher: all EP ranks must select the same CUDA graph. Sync batch dims
             # across the EP group so graph selection is consistent.
             adjusted_batch_dim = InferenceBatchDimensions.adjust_batch_dims_for_expert_parallelism(
-                real_batch_dim, ep_group=ep_group
+                real_batch_dim,
+                strict=strict,
+                decode_only_cuda_graphs=decode_only_cuda_graphs,
+                ep_group=ep_group,
+                smallest_non_decode_cuda_graph_size=smallest_non_decode_cuda_graph_size,
+                num_speculative_tokens=num_speculative_tokens,
+                ep_zmq_communicator=ep_zmq_communicator,
             )
+
             if adjusted_batch_dim is None:
+                # we hit this scenario if decode_only_cuda_graphs is true,
+                # and one of the EP ranks is running a non-decode step
+                # in that case, all ranks have to run in eager mode
                 return None
         else:
             adjusted_batch_dim = real_batch_dim