megatron-core 0.16.0rc0.dev101856__tar.gz → 0.16.0rc0.dev102154__tar.gz

{megatron_core-0.16.0rc0.dev101856 → megatron_core-0.16.0rc0.dev102154}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: megatron-core
-Version: 0.16.0rc0.dev101856
+Version: 0.16.0rc0.dev102154
 Summary: Megatron Core - a library for efficient and scalable training of transformer based models
 Author-email: NVIDIA <nemo-toolkit@nvidia.com>
 Maintainer-email: NVIDIA <nemo-toolkit@nvidia.com>

{megatron_core-0.16.0rc0.dev101856 → megatron_core-0.16.0rc0.dev102154}/megatron/core/fusions/fused_pad_routing_map.py

@@ -1,9 +1,11 @@
-# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
 from unittest.mock import MagicMock
 
 import torch
 from packaging import version
 
+from megatron.core.jit import jit_fuser
 from megatron.core.utils import experimental_fn, null_decorator
 
 try:
@@ -69,6 +71,7 @@ def _pad_routing_map_kernel(
 
 
 @experimental_fn(introduced_with_version="0.13.0")
+@jit_fuser
 def fused_pad_routing_map(routing_map: torch.Tensor, pad_multiple: int) -> torch.Tensor:
     """Fused version of pad_routing_map.
     Args:

{megatron_core-0.16.0rc0.dev101856 → megatron_core-0.16.0rc0.dev102154}/megatron/core/inference/contexts/attention_context/mamba_metadata.py

@@ -1,8 +1,28 @@
 # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
+from dataclasses import dataclass
+from typing import List, Optional, Tuple
+
 import torch
 
 
+@dataclass
+class MambaInferenceStateConfig:
+    """Config for initializing Mamba model inference state tensors."""
+
+    layer_type_list: List[str]
+    """
+    A list of strings that indicates the layer type (Mamba / Attention / MLP) for each layer.
+    See `megatron/core/ssm/mamba_hybrid_layer_allocation.py` for the list of symbols.
+    """
+
+    mamba_conv_states_shape: Tuple[int]
+    """Mamba conv states shape per request."""
+
+    mamba_ssm_states_shape: Tuple[int]
+    """Mamba ssm states shape per request."""
+
+
 class MambaMetadata:
     """Manages the metadata tensors required for Mamba layers during inference."""
 
@@ -64,7 +84,7 @@ class MambaMetadata:
         """
         self.request_to_mamba_state_idx_cudagraph_only[0:num_active_requests] = active_mamba_indices
 
-    def allocate_slot(self) -> int:
+    def allocate_slot(self) -> Optional[int]:
         """
         Allocates a new slot for a request in the Mamba state buffers.
 

{megatron_core-0.16.0rc0.dev101856 → megatron_core-0.16.0rc0.dev102154}/megatron/core/inference/contexts/dynamic_block_allocator.py

@@ -21,11 +21,11 @@ class BlockAllocator:
             space for paused requests that live on the CPU.
     """
 
-    def __init__(self, context: "DynamicInferenceContext", active_count: int):
+    def __init__(self, context: "DynamicInferenceContext", total_count: int):
 
         self.context = context
 
-        active_count -= 1  # -1 for dummy_block_idx (see below)
+        active_count = (total_count - 1) // 2  # -1 for dummy_block_idx (see below)
         active_count = max(1, active_count)  # need at least one block
         self.total_count = 2 * active_count + 1  # +1 for dummy_block_idx
         self.total_avail = self.total_count - 1  # -1 for dummy_block_idx

{megatron_core-0.16.0rc0.dev101856 → megatron_core-0.16.0rc0.dev102154}/megatron/core/inference/contexts/dynamic_context.py

@@ -24,14 +24,11 @@ from megatron.core.inference.unified_memory import (
 from megatron.core.inference.utils import tensor_swap
 from megatron.core.models.common.embeddings.rope_utils import apply_rotary_pos_emb
 from megatron.core.package_info import __version__ as mcore_version
-from megatron.core.ssm.mamba_hybrid_layer_allocation import (
-    Symbols,
-    get_layer_maps_from_layer_type_list,
-)
+from megatron.core.ssm.mamba_hybrid_layer_allocation import get_layer_maps_from_layer_type_list
 from megatron.core.transformer import TransformerConfig
 from megatron.core.utils import divide as core_divide
 
-from .attention_context.mamba_metadata import MambaMetadata
+from .attention_context.mamba_metadata import MambaInferenceStateConfig, MambaMetadata
 from .attention_context.mha_metadata import GraphedMHAMetadata, NonGraphedMHAMetadata
 from .base_context import BaseInferenceContext
 from .dynamic_block_allocator import BlockAllocator
@@ -202,37 +199,23 @@ class DynamicInferenceContext(BaseInferenceContext):
     arbitrary sequence length may be added, paused, or removed from the context
     at any step. The only constraint is the maximum number of requests or tokens
     that the context is defined to support. For the block-level KV cache, a memory
-    buffer is allocated up front (size `2 * active_buffer_size_gb`), that is
+    buffer is allocated up front (size `buffer_size_gb` if `unified_memory_level`
+    == 0, or `2 * buffer_size_gb` if `unified_memory_level` == 1), that is
     divided into blocks and dynamically assigned to requests. At any given step,
     any unassigned blocks equate to unused space.
 
-    Additionally, a fraction of the memory buffer (`gtd_request_fraction`, i.e.,
-    the 'guaranteed' request fraction) is reserved for guaranteeing that a
-    minimum number of active requests may continue to generate tokens on any step.
-    The reason for this is that the context manages two pools of requests: 1)
-    active requests, and 2) paused requests. Paused requests are requests where
-    insufficient memory blocks remain for future assignment, and these requests
-    are set aside until enough memory blocks are available. Active requests are
-    requests that have sufficient memory blocks to proceed with their generations.
-
-    The situation can arise where all requests eventually become paused due to all
-    memory blocks being assigned. In this case, there are no active requests and
-    thus no progress can be made. To handle this case, a fraction of the memory
-    buffer is reserved that only allows active requests, and no paused requests.
-    This fraction must be carefully tuned, as it can have an order of magnitude
-    impact on overall latency.
-
     Args:
         params_dtype (torch.dtype): Dtype used for KV cache.
-        num_layers (int): Number of layers.
+        num_layers (int): Number of layers on this pipeline parallel rank.
         kv_channels (int): Hidden dimension per attention head.
         num_attention_heads (int): Number of attention heads.
         max_sequence_length (int): Max possible sequence length (prompt + output)
             that will occur.
-        active_buffer_size_gb (float): Buffer size reserved for active requests
-            that live on the GPU. The total buffer size (stored in unified memory)
-            is 2x this value, with the the other half of the buffer reserved for
-            paused requests that live on the CPU.
+        buffer_size_gb (float): Buffer size reserved on the GPU for the KV cache.
+            if `unified_memory_level` >= 1, then CPU memory is additionally
+            utilized, resulting in a total buffer size of `2 * buffer_size_gb`.
+            Regardless of total buffer size, the KV cache is conceptually divided
+            into 50% active requests and 50% paused requests.
         max_tokens (int): Max number of tokens to use for forward passes. This is
             primarily limited by prefill activation memory usage. (Defaults to
             16384).
@@ -245,14 +228,8 @@ class DynamicInferenceContext(BaseInferenceContext):
         materialize_only_last_token_logits (Optional[bool]): Whether to only
             materialize logits for the last token. This should be set to False
             if returning log probs.
-        layer_type_list (Optional[List[str]]): A list of strings that indicates
-            the layer type (Mamba / Attention / MLP) for each layer.
-            See `megatron/core/ssm/mamba_hybrid_layer_allocation.py` for the list
-            of symbols. This must be provided for hybrid models.
-        mamba_conv_states_shape: (Optional[Tuple[int]]): Mamba conv states shape per request.
-            This must be provided for hybrid models.
-        mamba_ssm_states_shape: (Optional[Tuple[int]]): Mamba ssm states shape per request.
-            This must be provided for hybrid models.
+        mamba_inference_state_config (Optional[MambaInferenceStateConfig]): The Mamba
+            inference state config if the model is a hybrid model.
         use_cuda_graphs_for_non_decode_steps (bool): If True, use cuda graphs for non-decode
             engine steps.
         unified_memory_level (Optional[int]): Set unified memory usage within the
@@ -279,7 +256,7 @@ class DynamicInferenceContext(BaseInferenceContext):
         kv_channels: int,
         num_attention_heads: int,
         max_sequence_length: int,
-        active_buffer_size_gb: float,
+        buffer_size_gb: float,
         max_tokens: int = DEFAULT_MAX_TOKENS,
         block_size_tokens: int = 256,
         tensor_model_parallel_size: Optional[int] = None,
@@ -288,9 +265,7 @@ class DynamicInferenceContext(BaseInferenceContext):
         qk_pos_emb_head_dim: Optional[int] = None,
         num_cuda_graphs: Optional[int] = None,
         materialize_only_last_token_logits: Optional[bool] = True,
-        layer_type_list: Optional[List[str]] = None,
-        mamba_conv_states_shape: Optional[Tuple[int]] = None,
-        mamba_ssm_states_shape: Optional[Tuple[int]] = None,
+        mamba_inference_state_config: Optional[MambaInferenceStateConfig] = None,
         use_cuda_graphs_for_non_decode_steps: bool = True,
         use_flashinfer_fused_rope: bool = False,
         unified_memory_level: Optional[int] = 1,
@@ -317,8 +292,10 @@ class DynamicInferenceContext(BaseInferenceContext):
         self.num_attention_heads_per_partition = core_divide(num_attention_heads, tp_size)
 
         # Mamba states.
-        self.is_hybrid_model = layer_type_list is not None and Symbols.MAMBA in layer_type_list
+        self.is_hybrid_model = mamba_inference_state_config is not None
         if self.is_hybrid_model:
+            mamba_conv_states_shape = mamba_inference_state_config.mamba_conv_states_shape
+            mamba_ssm_states_shape = mamba_inference_state_config.mamba_ssm_states_shape
             assert (
                 mamba_conv_states_shape is not None
             ), "`mamba_conv_states_shape` must be specified for hybrid models"
@@ -333,7 +310,7 @@ class DynamicInferenceContext(BaseInferenceContext):
             # corresponding attention layer index or Mamba layer index depending on the
             # layer type.
             attention_layer_map, mamba_layer_map, _, _ = get_layer_maps_from_layer_type_list(
-                layer_type_list
+                mamba_inference_state_config.layer_type_list
             )
             self.num_attention_layers = len(attention_layer_map)
             self.num_mamba_layers = len(mamba_layer_map)
@@ -359,7 +336,10 @@ class DynamicInferenceContext(BaseInferenceContext):
             #   one vector  c_t  (rank)  +  optional RoPE phase slice
             self.kv_reduced_dim = kv_lora_rank + qk_pos_emb_head_dim
             self.block_size_bytes = (
-                dtype_size_bytes * num_layers * self.block_size_tokens * self.kv_reduced_dim
+                dtype_size_bytes
+                * num_attention_layers
+                * self.block_size_tokens
+                * self.kv_reduced_dim
             )
         else:
             self.block_size_bytes = (
@@ -379,14 +359,29 @@ class DynamicInferenceContext(BaseInferenceContext):
             mamba_states_memory_per_request *= self.num_mamba_layers
             mamba_states_memory_per_request *= dtype_size_bytes
 
+        # Unified memory.
+        self.unified_memory_level = unified_memory_level
+        if unified_memory_level > 0:
+            try:
+                self.unified_memory_mempool = create_unified_mempool()
+            except UnifiedMemoryUnsupportedError:
+                if torch.distributed.get_rank() == 0:
+                    warnings.warn(
+                        "Unified memory requested but not available; defaulting to GPU memory."
+                    )
+                self.unified_memory_level = 0
+
         # Initialize block allocator.
-        active_buffer_size_bytes = int(active_buffer_size_gb * 1024**3)
-        active_block_count_total = active_buffer_size_bytes // (
+        buffer_size_bytes = int(buffer_size_gb * 1024**3)
+        block_count_total = buffer_size_bytes // (
             self.block_size_bytes + mamba_states_memory_per_request
         )
-        self.block_allocator = BlockAllocator(context=self, active_count=active_block_count_total)
-        del active_block_count_total  # use self.block_allocator.active_count
-        active_buffer_size_bytes = self.block_allocator.active_count * self.block_size_bytes
+        self.block_allocator = BlockAllocator(
+            context=self,
+            total_count=(
+                block_count_total if self.unified_memory_level == 0 else 2 * block_count_total
+            ),
+        )
 
         # Set max_total_requests, max_active_requests, max_tokens.
         self.max_total_requests = self.block_allocator.total_count - 1  # -1 for dummy block
@@ -408,18 +403,6 @@ class DynamicInferenceContext(BaseInferenceContext):
         self.params_dtype = params_dtype
         self.max_sequence_length = max_sequence_length
 
-        # Unified memory.
-        self.unified_memory_level = unified_memory_level
-        if unified_memory_level > 0:
-            try:
-                self.unified_memory_mempool = create_unified_mempool()
-            except UnifiedMemoryUnsupportedError:
-                if torch.distributed.get_rank() == 0:
-                    warnings.warn(
-                        "Unified memory requested but not available; defaulting to GPU memory."
-                    )
-                self.unified_memory_level = 0
-
         # Request and token counts.
         self.total_request_count = 0
         self.active_token_count = 0
@@ -513,7 +496,10 @@ class DynamicInferenceContext(BaseInferenceContext):
         # Print info.
         logging.info(
             "DynamicInferenceContext: allocated context with active buffer size %s (%d blocks)."
-            % (get_mem_size_str(active_buffer_size_bytes), self.block_allocator.active_count)
+            % (
+                get_mem_size_str(self.block_allocator.active_count * self.block_size_bytes),
+                self.block_allocator.active_count,
+            )
         )
 
     def allocate_all_tensors(self, *, is_init: bool) -> None:
@@ -731,8 +717,9 @@ class DynamicInferenceContext(BaseInferenceContext):
         inference_config: InferenceWrapperConfig,
         model,
         max_batch_size: int,
-        active_buffer_size_gb: float = 40,
+        buffer_size_gb: float = 40,
         num_cuda_graphs: int = None,
+        mamba_inference_state_config: Optional[MambaInferenceStateConfig] = None,
     ):
         """
         Instantiate a `DynamicInferenceContext` from a `TransformerConfig` and an `InferenceWrapperConfig`.
@@ -746,14 +733,15 @@ class DynamicInferenceContext(BaseInferenceContext):
         max_sequence_length = max(max_sequence_length, max_batch_size)
         return cls(
             params_dtype=inference_config.params_dtype,
-            num_layers=model_config.num_layers,
+            num_layers=model_config.num_layers // model_config.pipeline_model_parallel_size,
             kv_channels=model_config.kv_channels,
             num_attention_heads=model_config.num_query_groups,
             max_sequence_length=inference_config.inference_max_seq_length,
-            active_buffer_size_gb=active_buffer_size_gb,
+            buffer_size_gb=buffer_size_gb,
             materialize_only_last_token_logits=False,
             num_cuda_graphs=num_cuda_graphs,
             use_flashinfer_fused_rope=None,
+            mamba_inference_state_config=mamba_inference_state_config,
         )
 
     @classmethod
@@ -894,6 +882,7 @@ class DynamicInferenceContext(BaseInferenceContext):
             to blocks within the block-level memory buffer.
         """
         attention_layer_number = self.layer_map[layer_number - 1]
+
         if self.cache_mla_latent:
             return (
                 self.memory_buffer[attention_layer_number],

{megatron_core-0.16.0rc0.dev101856 → megatron_core-0.16.0rc0.dev102154}/megatron/core/inference/engines/dynamic_engine.py

@@ -4,11 +4,13 @@ import asyncio
 import logging
 import multiprocessing
 import os
+import socket
 import struct
 import time
 import warnings
 from collections import deque
 from contextlib import contextmanager
+from dataclasses import dataclass
 from datetime import datetime
 from itertools import repeat
 from typing import Dict, List, Optional, Tuple, Union
@@ -94,6 +96,14 @@ def format_mem_bytes(mem_bytes):
     return "%d bytes" % mem_bytes
 
 
+@dataclass(kw_only=True)
+class RequestEntry:
+    """Entry in the engine's `self.requests` dict."""
+
+    record: DynamicInferenceRequestRecord
+    future: asyncio.Future
+
+
 # pylint: disable=line-too-long
 class DynamicInferenceEngine(AbstractEngine):
     """The dynamic inference engine.
@@ -214,8 +224,7 @@ class DynamicInferenceEngine(AbstractEngine):
         self.waiting_request_ids = deque()
         self.failed_request_ids = []  # deque()
         self.request_counter = Counter()
-        self.request_records: Dict[int, DynamicInferenceRequestRecord] = {}
-        self.request_completion_futures: Dict[int, asyncio.Future] = {}
+        self.requests: Dict[int, RequestEntry] = {}
         self.step_start_event = torch.cuda.Event(enable_timing=True)
         self.step_end_event = torch.cuda.Event(enable_timing=True)
         self.paused = False
@@ -338,6 +347,7 @@ class DynamicInferenceEngine(AbstractEngine):
     async def start_listening_to_data_parallel_coordinator(
         self,
         inference_coordinator_port: int,
+        inference_mp_coordinator_port: int = 20000,
         launch_inference_coordinator: bool = True,
         verbose: bool = False,
         *,
@@ -351,15 +361,15 @@ class DynamicInferenceEngine(AbstractEngine):
         based on the rank's role within the distributed topology.
 
         The setup involves two primary roles within each data-parallel group:
-        1.  **TP Coordinator (TP_rank=0, PP_rank=0)**: This rank connects directly
+        1.  **MP Coordinator (TP_rank=0, PP_rank=0)**: This rank connects directly
             to the central coordinator via a ZMQ `DEALER` socket. It receives
             requests and uses a ZMQ `PUB` (publisher) socket to broadcast them
-            to all other ranks within its tensor-parallel (TP) group.
-        2.  **TP Workers (all other ranks)**: These ranks use ZMQ `SUB` (subscriber)
-            sockets to listen for requests broadcast by their local TP Coordinator.
+            to all other ranks within its model-parallel (MP) group.
+        2.  **MP Workers (all other ranks)**: These ranks use ZMQ `SUB` (subscriber)
+            sockets to listen for requests broadcast by their local MP Coordinator.
 
-        This architecture uses fast Inter-Process Communication (`ipc`) sockets for
-        intra-node broadcasts within a TP group.
+        This architecture uses TCP sockets for both inter-node and intra-node broadcasts
+        within an MP group.
 
         Finally, after setting up the communication channels and ensuring all ranks
         are synchronized, this method starts the main engine processing loop
@@ -368,16 +378,13 @@ class DynamicInferenceEngine(AbstractEngine):
         Args:
             inference_coordinator_port (int): The network port where the central
                 `InferenceCoordinator` is or will be listening.
+            inference_mp_coordinator_port (int): The base network port where each model parallel
+                coordinator will broadcast messages from. Each MP group will compute an independent
+                port offset from this base port.
             launch_inference_coordinator (bool, optional): If True, the global rank 0
                 process will spawn and manage the `InferenceCoordinator`
                 process. Defaults to True.
             verbose (bool): Whether to run in verbose mode.
-
-        Note:
-            The current implementation uses `ipc` sockets for broadcasting requests
-            within a Tensor Parallel group, which limits each TP group to a single
-            physical node. For example, if you have 8 GPUs per node, then this will only
-            work with TP=[1,2,4,8]
         """
 
         assert HAVE_ZMQ, (
@@ -404,8 +411,31 @@ class DynamicInferenceEngine(AbstractEngine):
         # Todo [Siddharth]: can we move this code to another file?
         self.zmq_context = zmq.Context()
         self.zmq_sockets = []  # keep track of all sockets created by this engine
+
+        # We need to broadcast the hostname of the (TP=0, PP=0) rank
+        # to all other ranks in the same model parallel group.
+        tp_rank = parallel_state.get_tensor_model_parallel_rank()
+        pp_rank = parallel_state.get_pipeline_model_parallel_rank()
+
+        hostname_list = [None]
+        if tp_rank == 0 and pp_rank == 0:
+            hostname_list[0] = socket.gethostname()
+
+        # Find the global rank of the (TP=0, PP=0) rank in our MP group
+        src_global_rank = parallel_state.get_model_parallel_src_rank()
+
+        torch.distributed.broadcast_object_list(
+            hostname_list, src=src_global_rank, group=parallel_state.get_model_parallel_group()
+        )
+        bcast_hostname = hostname_list[0]
+
+        # We need unique ports for each MP group, so we compute an offset using the DP rank.
+        dp_rank = parallel_state.get_data_parallel_rank()
+        req_port = inference_mp_coordinator_port + (dp_rank * 2)
+        len_port = inference_mp_coordinator_port + (dp_rank * 2) + 1
+
         ip_address_of_dp_coordinator = os.getenv('MASTER_ADDR', '127.0.0.1')
-        identity = f'tp-coord-{parallel_state.get_data_parallel_rank()}'
+        identity = f'mp-coord-{parallel_state.get_data_parallel_rank()}'
         if (
             parallel_state.get_tensor_model_parallel_rank() == 0
             and parallel_state.get_pipeline_model_parallel_rank() == 0
@@ -423,37 +453,33 @@ class DynamicInferenceEngine(AbstractEngine):
             self.socket_for_receiving_requests.send(b"")
 
             # 2. Create a publisher socket. This is used to publish or broadcast
-            #    requests within the tensor parallel group
-            self.tensor_parallel_publisher_socket = self.zmq_context.socket(zmq.PUB)
-            self.tensor_parallel_publisher_socket.bind(f"ipc:///tmp/{identity}-tp-bcast-socket-req")
+            #    requests within the model parallel group
+            self.model_parallel_publisher_socket = self.zmq_context.socket(zmq.PUB)
+            self.model_parallel_publisher_socket.bind(f"tcp://*:{req_port}")
 
             # 3. Create another publisher socket to broadcast the number of messages to receive.
-            self.tensor_parallel_num_msgs_publisher_socket = self.zmq_context.socket(zmq.PUB)
-            self.tensor_parallel_num_msgs_publisher_socket.bind(
-                f"ipc:///tmp/{identity}-tp-bcast-socket-len"
-            )
+            self.model_parallel_num_msgs_publisher_socket = self.zmq_context.socket(zmq.PUB)
+            self.model_parallel_num_msgs_publisher_socket.bind(f"tcp://*:{len_port}")
             self.zmq_sockets += [
                 self.socket_for_receiving_requests,
-                self.tensor_parallel_num_msgs_publisher_socket,
-                self.tensor_parallel_publisher_socket,
+                self.model_parallel_num_msgs_publisher_socket,
+                self.model_parallel_publisher_socket,
             ]
-        # All TP ranks subscribe to the two publisher sockets
-        self.tensor_parallel_subscriber_socket = self.zmq_context.socket(zmq.SUB)
-        self.tensor_parallel_subscriber_socket.connect(f"ipc:///tmp/{identity}-tp-bcast-socket-req")
-        self.tensor_parallel_subscriber_socket.setsockopt_string(zmq.SUBSCRIBE, "")
-
-        self.tensor_parallel_num_msgs_subscriber_socket = self.zmq_context.socket(zmq.SUB)
-        self.tensor_parallel_num_msgs_subscriber_socket.connect(
-            f"ipc:///tmp/{identity}-tp-bcast-socket-len"
-        )
-        self.tensor_parallel_num_msgs_subscriber_socket.setsockopt_string(zmq.SUBSCRIBE, "")
+        # All MP ranks subscribe to the two publisher sockets
+        self.model_parallel_subscriber_socket = self.zmq_context.socket(zmq.SUB)
+        self.model_parallel_subscriber_socket.connect(f"tcp://{bcast_hostname}:{req_port}")
+        self.model_parallel_subscriber_socket.setsockopt_string(zmq.SUBSCRIBE, "")
+
+        self.model_parallel_num_msgs_subscriber_socket = self.zmq_context.socket(zmq.SUB)
+        self.model_parallel_num_msgs_subscriber_socket.connect(f"tcp://{bcast_hostname}:{len_port}")
+        self.model_parallel_num_msgs_subscriber_socket.setsockopt_string(zmq.SUBSCRIBE, "")
 
         self.zmq_sockets += [
-            self.tensor_parallel_subscriber_socket,
-            self.tensor_parallel_num_msgs_subscriber_socket,
+            self.model_parallel_subscriber_socket,
+            self.model_parallel_num_msgs_subscriber_socket,
         ]
 
-        torch.distributed.barrier(parallel_state.get_tensor_model_parallel_group())
+        torch.distributed.barrier(parallel_state.get_model_parallel_group())
 
         if launch_inference_coordinator and torch.distributed.get_rank() == 0:
             await await_process_event(coordinator_ready_event, self.inference_coordinator_process)
@@ -549,13 +575,13 @@ class DynamicInferenceEngine(AbstractEngine):
 
         # Maintain references to requests before reset.
         waiting_request_ids = list(self.waiting_request_ids)
-        active_request_ids = set(self.request_records.keys()) - set(waiting_request_ids)
+        active_request_ids = set(self.requests.keys()) - set(waiting_request_ids)
         self.resume_request_ids = [*active_request_ids, *waiting_request_ids]
         self.waiting_request_ids.clear()
 
         # Suspend requests objects.
         for request_id in active_request_ids:
-            self.request_records[request_id].suspend(self.controller.tokenizer)
+            self.requests[request_id].record.suspend(self.controller.tokenizer)
 
     def resume(self):
         """Resume engine by reallocating context's GPU state."""
@@ -632,7 +658,7 @@ class DynamicInferenceEngine(AbstractEngine):
         Returns:
             (DynamicInferenceRequest) The most recent request in the record.
         """
-        return self.request_records[request_id][-1]
+        return self.requests[request_id].record[-1]
 
     def _add_request(
         self, request: DynamicInferenceRequest
@@ -640,10 +666,13 @@ class DynamicInferenceEngine(AbstractEngine):
 
         request_id = request.request_id
 
-        # Add request to self.request_records. If the engine has previously been
+        # Add request to self.requests. If the engine has previously been
         # suspended, then the request may already exist.
-        if request_id not in self.request_records:
-            self.request_records[request_id] = DynamicInferenceRequestRecord.from_request(request)
+        if request_id not in self.requests:
+            self.requests[request_id] = RequestEntry(
+                record=DynamicInferenceRequestRecord.from_request(request),
+                future=self._loop.create_future(),
+            )
 
         if request.status is None:
             request.status = Status.ACTIVE_AND_GENERATING_TOKENS
@@ -686,13 +715,7 @@ class DynamicInferenceEngine(AbstractEngine):
         if request.status != Status.FAILED:
             self.waiting_request_ids.append(request_id)
 
-        # Create a new asyncio Future to notify the user when the request has
-        # completed. If the engine has previously been suspended, then the future
-        # may already exist.
-        if request_id not in self.request_completion_futures:
-            self.request_completion_futures[request_id] = self._loop.create_future()
-
-        return self.request_completion_futures[request_id]
+        return self.requests[request_id].future
 
     def add_request(
         self,
@@ -813,8 +836,8 @@ class DynamicInferenceEngine(AbstractEngine):
                 if request_id in finished_request_ids:
                     request.generated_length = len(request.generated_tokens)
                     request.status = Status.COMPLETED
-                    finished_request_record = self.request_records.pop(request_id)
-                    finished_request = finished_request_record[-1]
+                    finished_entry = self.requests.pop(request_id)
+                    finished_request = finished_entry.record[-1]
                     if finished_request.prompt is None:
                         finished_request.prompt = self.controller.tokenizer.detokenize(
                             finished_request.prompt_tokens.tolist()
@@ -823,8 +846,8 @@ class DynamicInferenceEngine(AbstractEngine):
                     finished_request.generated_text = self.controller.tokenizer.detokenize(
                         finished_request.generated_tokens
                     )
-                    finished_request_records.append(finished_request_record)
-                    self.request_completion_futures[request_id].set_result(finished_request_record)
+                    finished_request_records.append(finished_entry.record)
+                    finished_entry.future.set_result(finished_entry.record)
                 else:
                     active_request_ids.append(request_id)
             else:
@@ -1007,12 +1030,12 @@ class DynamicInferenceEngine(AbstractEngine):
 
         # Failed requests.
         for failed_request_id in self.failed_request_ids:
-            failed_request_record = self.request_records.pop(failed_request_id)
-            failed_request = failed_request_record[-1]
+            failed_entry = self.requests.pop(failed_request_id)
+            failed_request = failed_entry.record[-1]
             failed_request.status = Status.FAILED
             failed_request.add_event_fail()
-            finished_request_records.append(failed_request_record)
-            self.request_completion_futures[failed_request_id].set_result(failed_request_record)
+            finished_request_records.append(failed_entry.record)
+            failed_entry.future.set_result(failed_entry.record)
         self.failed_request_ids.clear()
 
         # Log KV cache utilization stats to W&B
@@ -1154,16 +1177,16 @@ class DynamicInferenceEngine(AbstractEngine):
         """Drains the ZMQ socket for a batch of requests and adds them to the engine.
 
         This method is a collective and synchronous operation that must be called
-        by all ranks in a Tensor Parallel (TP) group at the same time. It ensures
+        by all ranks in a Model Parallel (MP) group at the same time. It ensures
         that all ranks process the exact same batch of incoming requests and
         control signals.
 
         The synchronization works as follows:
-        1.  The TP rank 0 drains all pending messages from its subscriber socket
+        1.  The MP rank 0 drains all pending messages from its subscriber socket
             in a non-blocking manner.
-        2.  TP rank 0 then broadcasts the number of messages it received to all other
-            ranks in its TP group using a dedicated publisher socket.
-        3.  The other TP ranks wait to receive this count, and then receive exactly
+        2.  MP rank 0 then broadcasts the number of messages it received to all other
+            ranks in its MP group using a dedicated publisher socket.
+        3.  The other MP ranks wait to receive this count, and then receive exactly
             that many messages from their subscriber sockets.
 
         Once all ranks have the same batch of messages, they are unpacked and
@@ -1173,7 +1196,7 @@ class DynamicInferenceEngine(AbstractEngine):
 
         Note:
             This function is synchronous and must be called collectively by all
-            ranks in a TP group. It should not be launched in a separate coroutine
+            ranks in a MP group. It should not be launched in a separate coroutine
             to ensure all ranks execute it in lockstep before proceeding to the
             next engine step.
 
@@ -1181,10 +1204,11 @@ class DynamicInferenceEngine(AbstractEngine):
             int: The number of messages that were received and processed in this batch.
         """
 
-        rank = parallel_state.get_tensor_model_parallel_rank()
+        tp_rank = parallel_state.get_tensor_model_parallel_rank()
+        pp_rank = parallel_state.get_pipeline_model_parallel_rank()
         torch.cuda.nvtx.range_push("drain_zmq_socket")
         all_messages = []
-        if rank == 0:
+        if tp_rank == 0 and pp_rank == 0:
             while True:
                 try:
                     # Receive messages in a non-blocking way.
@@ -1196,22 +1220,22 @@ class DynamicInferenceEngine(AbstractEngine):
             # First publish the number of messages to dequeue.
             # This is important because we want all tensor parallel ranks
             # to dequeue the same number of messages.
-            self.tensor_parallel_num_msgs_publisher_socket.send(
+            self.model_parallel_num_msgs_publisher_socket.send(
                 struct.pack('!i', messages_to_dequeue)
             )
-            # Now publish the actual messages to all tensor parallel ranks
+            # Now publish the actual messages to all model parallel ranks
             for message in all_messages:
-                self.tensor_parallel_publisher_socket.send(message)
+                self.model_parallel_publisher_socket.send(message)
         else:
-            # First, receive the number of messages to dequeue from tp-rank 0
+            # First, receive the number of messages to dequeue from mp-rank 0
             messages_to_dequeue = struct.unpack(
-                '!i', self.tensor_parallel_num_msgs_subscriber_socket.recv()
+                '!i', self.model_parallel_num_msgs_subscriber_socket.recv()
             )[0]
             # Now, dequeue the same number of messages from the subscriber socket.
             # Note that these receives are blocking, because the messages
             # are guaranteed to be available after the tp-rank 0 has sent them.
             for _ in range(messages_to_dequeue):
-                all_messages.append(self.tensor_parallel_subscriber_socket.recv())
+                all_messages.append(self.model_parallel_subscriber_socket.recv())
 
         torch.cuda.nvtx.range_pop()
         for message in all_messages: