diff --git a/docs/user-guide/features/megatron_fsdp.md b/docs/user-guide/features/megatron_fsdp.md index 6b55bf9fcce..36fcc68893c 100644 --- a/docs/user-guide/features/megatron_fsdp.md +++ b/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). \ No newline at end of file +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). diff --git a/megatron/core/inference/batch_dimensions_utils.py b/megatron/core/inference/batch_dimensions_utils.py index 223368b4084..6ccbccf0f33 100644 --- a/megatron/core/inference/batch_dimensions_utils.py +++ b/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 diff --git a/megatron/core/inference/contexts/attention_context/mamba_metadata.py b/megatron/core/inference/contexts/attention_context/mamba_metadata.py index b07802a8f8c..ff6423be16b 100644 --- a/megatron/core/inference/contexts/attention_context/mamba_metadata.py +++ b/megatron/core/inference/contexts/attention_context/mamba_metadata.py @@ -35,9 +35,9 @@ def __init__( # Maximum possible chunks across all batch configurations self.max_chunks = max_tokens // mamba_chunk_size + max_requests - # Map from requests to slots in the static Mamba state buffer + # Map from requests to slots in the static Mamba state buffer (CPU for bookkeeping). self.request_to_mamba_state_idx = torch.full( - (self.max_requests,), -1, dtype=torch.int32, device=torch.cuda.current_device() + (self.max_requests,), -1, dtype=torch.int32, device='cpu' ) # Map from requests to slots in the static Mamba state buffer for active decode requests. @@ -85,9 +85,9 @@ def __init__( self._conv_seq_idx_buffer = torch.zeros(max_tokens, dtype=torch.int32, device=self.device) self._conv_seq_start_buffer = torch.zeros(max_tokens, dtype=torch.int32, device=self.device) - # Allocator for Mamba state slots + # Allocator for Mamba state slots (CPU for bookkeeping). self.mamba_state_free_slots = torch.arange( - self.max_requests, dtype=torch.int32, device=torch.cuda.current_device() + self.max_requests, dtype=torch.int32, device='cpu' ) self.mamba_state_free_slot_count = self.max_requests @@ -108,8 +108,31 @@ def __init__( else: self.conv_gather_offsets = None + # Coalesced production path: pinned CPU views + shared GPU views bound + # by DynamicInferenceContext so that the per-step Mamba metadata fields + # ride along with the single coalesced H2D in transfer_bookkeeping_to_gpu. + # The legacy update() path above keeps using the standalone _*_buffer + # tensors (exercised only by unit tests that construct MambaMetadata + # without a context). + self._cpu_bufs = None + self._gpu_view = None + self.reset_varlen_metadata() + def bind_cpu_buffers(self, bufs: dict) -> None: + """Attach pinned CPU views from DynamicInferenceContext._cpu_bookkeeping_buf. + + ``bufs`` maps field names to 1D (or (1, max_tokens) for ``seq_idx``) + pinned CPU views that compute_cpu_metadata writes into. The matching + GPU views on the other side of the H2D are exposed via + :meth:`bind_gpu_buffers`. + """ + self._cpu_bufs = bufs + + def bind_gpu_buffers(self, gpu_view) -> None: + """Attach shared GPU views from the context's :class:`ContextGPUView`.""" + self._gpu_view = gpu_view + def reset(self) -> None: """ Resets all Mamba states and frees all allocated slots. @@ -120,7 +143,7 @@ def reset(self) -> None: # Re-initialize the free slot pool self.mamba_state_free_slots = torch.arange( - self.max_requests, dtype=torch.int32, device=torch.cuda.current_device() + self.max_requests, dtype=torch.int32, device='cpu' ) self.mamba_state_free_slot_count = self.max_requests @@ -340,6 +363,7 @@ def _update_intermediate_metadata( intermediate_offsets_gpu: Optional[torch.Tensor], intermediate_counts_gpu: Optional[torch.Tensor], real_prefill_count: int, + cu_seqlens_gpu: Optional[torch.Tensor] = None, ) -> None: """Precompute intermediate extraction metadata for CUDA graph compatibility. @@ -353,18 +377,32 @@ def _update_intermediate_metadata( intermediate_counts_gpu: [real_prefill_count] int32 GPU tensor of per-request offset counts (0-3), or None. real_prefill_count: Number of real (non-padding) prefill requests. + cu_seqlens_gpu: GPU cu_seqlens tensor to read from. Defaults to + the legacy standalone ``_cu_seqlens_buffer`` used by + :meth:`update`; the coalesced production path passes the + shared ``ContextGPUView.mamba_cu_seqlens`` view. """ chunk_size = self.mamba_chunk_size max_count = self.max_intermediate_count + if cu_seqlens_gpu is None: + cu_seqlens_gpu = self._cu_seqlens_buffer if intermediate_offsets_gpu is not None and real_prefill_count > 0: - # Transfer counts to CPU (single sync) for per_request_counts and total check + # counts_list is CPU-cheap (source is already CPU from MambaSlotAllocator). counts_list = intermediate_counts_gpu.tolist() total = sum(counts_list) + # Ensure GPU copies for vectorized GPU ops below. + if not intermediate_offsets_gpu.is_cuda: + intermediate_offsets_gpu = intermediate_offsets_gpu.to( + self.device, non_blocking=True + ) + if not intermediate_counts_gpu.is_cuda: + intermediate_counts_gpu = intermediate_counts_gpu.to(self.device, non_blocking=True) + if total > 0: # Compute cumulative chunk counts from cu_seqlens (already on GPU) - cu = self._cu_seqlens_buffer[: real_prefill_count + 1] + cu = cu_seqlens_gpu[: real_prefill_count + 1] seq_lens = (cu[1 : real_prefill_count + 1] - cu[:real_prefill_count]).to( torch.int64 ) @@ -430,6 +468,223 @@ def _update_intermediate_metadata( self.intermediate_chunk_indices = self._intermediate_chunk_indices_buffer[:max_count] self.intermediate_abs_positions = self._intermediate_abs_positions_buffer[:max_count] + def compute_cpu_metadata( + self, + active_mamba_indices: torch.Tensor, + token_to_request_idx: torch.Tensor, + cpu_cu_query: torch.Tensor, + batch_dimensions: InferenceBatchDimensions, + padded_batch_dimensions: InferenceBatchDimensions, + enable_chunked_prefill: bool, + intermediate_offsets_gpu: Optional[torch.Tensor] = None, + intermediate_counts_gpu: Optional[torch.Tensor] = None, + ) -> dict: + """Compute all Mamba metadata on CPU, writing directly into the bound + pinned CPU views. + + The values written here are transferred to GPU by the single coalesced + H2D in :meth:`DynamicInferenceContext.transfer_bookkeeping_to_gpu`. + The returned dict contains only Python scalars + the intermediate GPU + tensors, which :meth:`load_from_cpu` consumes after the H2D. + + Args: + active_mamba_indices: CPU tensor of Mamba slot indices for active requests. + token_to_request_idx: CPU tensor mapping tokens to request indices. + cpu_cu_query: CPU cumulative query lengths from MHA metadata computation. + batch_dimensions: Dimensions of the current batch. + padded_batch_dimensions: Dimensions of the padded batch. + enable_chunked_prefill: Whether chunked prefill is enabled. + intermediate_offsets_gpu: GPU tensor of per-request intermediate offsets, or None. + intermediate_counts_gpu: GPU tensor of per-request intermediate counts, or None. + """ + assert self._cpu_bufs is not None, "bind_cpu_buffers() must be called first" + bufs = self._cpu_bufs + + real_decode_count = batch_dimensions.decode_req_count + real_prefill_count = batch_dimensions.prefill_req_count + 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 + chunk_size = self.mamba_chunk_size + + result = { + "padded_decode_count": padded_decode_count, + "padded_prefill_count": padded_prefill_count, + "padded_token_count": padded_token_count, + "real_decode_count": real_decode_count, + "real_prefill_count": real_prefill_count, + } + + # Decode batch indices (write into pinned view; padded slots = -1). + if padded_decode_count > 0: + bufs['batch_indices_decode'][:real_decode_count] = active_mamba_indices[ + :real_decode_count + ] + if padded_decode_count > real_decode_count: + bufs['batch_indices_decode'][real_decode_count:padded_decode_count] = -1 + + # Prefill batch indices, seq_idx, cu_seqlens, chunk/conv metadata. + if padded_prefill_count > 0: + if real_prefill_count > 0: + start = real_decode_count + bufs['batch_indices_prefill'][:real_prefill_count] = active_mamba_indices[ + start : start + real_prefill_count + ] + if padded_prefill_count > real_prefill_count: + bufs['batch_indices_prefill'][real_prefill_count:padded_prefill_count] = -1 + + # seq_idx: normalized token-to-request mapping for prefill tokens. + prefill_start_req = real_decode_count + end_prefill_req = real_decode_count + real_prefill_count + start_token = cpu_cu_query[prefill_start_req].item() + end_token = cpu_cu_query[end_prefill_req].item() + seq_len = end_token - start_token + + if seq_len > 0: + raw = token_to_request_idx[start_token:end_token] + bufs['seq_idx'][0, :seq_len] = raw - raw[0] + if padded_token_count > seq_len: + bufs['seq_idx'][0, seq_len:padded_token_count] = -1 + result["seq_len"] = seq_len + + # cu_seqlens for prefill. + cu_seqlens_view = bufs['cu_seqlens'] + cu_seqlens_view[0] = 0 + if real_prefill_count > 0: + cu_seqlens_view[1 : real_prefill_count + 1] = ( + cpu_cu_query[prefill_start_req + 1 : end_prefill_req + 1] + - cpu_cu_query[prefill_start_req] + ) + if real_prefill_count < padded_prefill_count: + last_val = cu_seqlens_view[real_prefill_count].item() + cu_seqlens_view[real_prefill_count + 1 : padded_prefill_count + 1] = last_val + + cu_seqlens_list = cu_seqlens_view[: real_prefill_count + 1].tolist() + real_prefill_tokens = ( + cu_seqlens_list[real_prefill_count] if real_prefill_count > 0 else 0 + ) + result["cu_seqlens_list"] = cu_seqlens_list + result["real_prefill_token_count"] = real_prefill_tokens + + # Chunk metadata (Python loop, pure CPU). + cu_seqlens_all = cu_seqlens_view[: padded_prefill_count + 1].tolist() + chunk_boundaries = [0] + last_chunk_idx_list = [] + chunk_to_seq_list = [] + + for i in range(padded_prefill_count): + start = cu_seqlens_all[i] + end = cu_seqlens_all[i + 1] + s_len = end - start + n_chunks = max(1, (s_len + chunk_size - 1) // chunk_size) + boundaries = [min(start + (k + 1) * chunk_size, end) for k in range(n_chunks)] + chunk_boundaries.extend(boundaries) + chunk_to_seq_list.extend([i] * n_chunks) + last_chunk_idx_list.append(len(chunk_boundaries) - 2) + + padded_max_chunks = padded_token_count // chunk_size + padded_prefill_count + last_boundary = chunk_boundaries[-1] + pad_b = padded_max_chunks + 1 - len(chunk_boundaries) + if pad_b > 0: + chunk_boundaries.extend([last_boundary] * pad_b) + pad_s = padded_max_chunks - len(chunk_to_seq_list) + if pad_s > 0: + chunk_to_seq_list.extend([0] * pad_s) + + n_cu = padded_max_chunks + 1 + bufs['cu_chunk_seqlens'][:n_cu] = torch.tensor( + chunk_boundaries[:n_cu], dtype=torch.int32 + ) + bufs['last_chunk_indices'][:padded_prefill_count] = torch.tensor( + last_chunk_idx_list, dtype=torch.int32 + ) + bufs['seq_idx_for_varlen'][:padded_max_chunks] = torch.tensor( + chunk_to_seq_list[:padded_max_chunks], dtype=torch.int32 + ) + result["padded_max_chunks"] = padded_max_chunks + + # Conv1d per-token metadata (CPU repeat_interleave). + conv_seq_idx_view = bufs['conv_seq_idx'] + conv_seq_start_view = bufs['conv_seq_start'] + if real_prefill_tokens > 0: + cu_t = cu_seqlens_view[: real_prefill_count + 1] + lengths = (cu_t[1:] - cu_t[:-1]).to(torch.int64) + seq_indices = torch.arange(real_prefill_count, dtype=torch.int32) + seq_starts = cu_t[:real_prefill_count].to(torch.int32) + conv_seq_idx_view[:real_prefill_tokens] = torch.repeat_interleave( + seq_indices, lengths + ) + conv_seq_start_view[:real_prefill_tokens] = torch.repeat_interleave( + seq_starts, lengths + ) + if padded_token_count > real_prefill_tokens: + conv_seq_idx_view[real_prefill_tokens:padded_token_count] = 0 + conv_seq_start_view[real_prefill_tokens:padded_token_count] = 0 + + # Intermediate metadata still requires GPU data: defer to load_from_cpu. + result["intermediate_offsets_gpu"] = intermediate_offsets_gpu + result["intermediate_counts_gpu"] = intermediate_counts_gpu + + # device_decode_prefill scalars. + if padded_decode_count > 0 and padded_prefill_count > 0: + result["decode_prefill_0"] = cpu_cu_query[real_decode_count].item() + result["decode_prefill_1"] = ( + cpu_cu_query[real_decode_count + real_prefill_count].item() + - cpu_cu_query[real_decode_count].item() + ) + + return result + + def load_from_cpu(self, d: dict) -> None: + """Point state attributes at the freshly-transferred shared GPU views. + + No H2D copies happen here: the Mamba metadata fields were transferred + as part of the coalesced bookkeeping H2D. This method just slices the + bound GPU views to the per-step sizes and runs the intermediate + metadata computation (which reads from the now-valid GPU cu_seqlens). + + Args: + d: Dict returned by compute_cpu_metadata(). + """ + assert self._gpu_view is not None, "bind_gpu_buffers() must be called first" + v = self._gpu_view + + padded_decode_count = d["padded_decode_count"] + padded_prefill_count = d["padded_prefill_count"] + padded_token_count = d["padded_token_count"] + real_prefill_count = d["real_prefill_count"] + + if padded_decode_count > 0: + self.batch_indices_decode = v.mamba_batch_indices_decode[:padded_decode_count] + + if padded_prefill_count > 0: + self.batch_indices_prefill = v.mamba_batch_indices_prefill[:padded_prefill_count] + self.seq_idx = v.mamba_seq_idx[:, :padded_token_count] + self.cu_seqlens = v.mamba_cu_seqlens[: padded_prefill_count + 1] + self.cu_seqlens_list = d["cu_seqlens_list"] + self.real_prefill_token_count = d["real_prefill_token_count"] + + padded_max_chunks = d["padded_max_chunks"] + self.cu_chunk_seqlens = v.mamba_cu_chunk_seqlens[: padded_max_chunks + 1] + self.last_chunk_indices = v.mamba_last_chunk_indices[:padded_prefill_count] + self.seq_idx_for_varlen = v.mamba_seq_idx_for_varlen[:padded_max_chunks] + self.conv_seq_idx = v.mamba_conv_seq_idx[:padded_token_count] + self.conv_seq_start = v.mamba_conv_seq_start[:padded_token_count] + + # Intermediate metadata reads from the just-transferred cu_seqlens + # to compute chunk indices & absolute positions for state extraction. + self._update_intermediate_metadata( + d["intermediate_offsets_gpu"], + d["intermediate_counts_gpu"], + real_prefill_count, + cu_seqlens_gpu=v.mamba_cu_seqlens, + ) + + if padded_decode_count > 0 and padded_prefill_count > 0: + self._device_decode_prefill_buffer[0] = d["decode_prefill_0"] + self._device_decode_prefill_buffer[1] = d["decode_prefill_1"] + self.device_decode_prefill = self._device_decode_prefill_buffer + def allocate_slot(self) -> Optional[int]: """ Allocates a new slot for a request in the Mamba state buffers. diff --git a/megatron/core/inference/contexts/attention_context/mha_metadata.py b/megatron/core/inference/contexts/attention_context/mha_metadata.py index 07f8a349b51..a71da895ea5 100644 --- a/megatron/core/inference/contexts/attention_context/mha_metadata.py +++ b/megatron/core/inference/contexts/attention_context/mha_metadata.py @@ -1,215 +1,84 @@ # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. import torch -from megatron.core.inference.batch_dimensions_utils import InferenceBatchDimensions - from .metadata_base import MetadataBase class MHAMetadata(MetadataBase): """ Metadata for MHA layer using flash-attention. + + GPU storage for the per-step fields (``query_lengths``, + ``cu_query_seq_lengths``, ``kv_seq_lengths``, ``cu_kv_seq_lengths``, + ``block_table``) lives inside the context's :class:`ContextGPUView` + unified buffer. Both :class:`GraphedMHAMetadata` and + :class:`NonGraphedMHAMetadata` bind to the same GPU views (only one is + active per step), so the single coalesced H2D in + :meth:`DynamicInferenceContext.transfer_bookkeeping_to_gpu` covers the + MHA fields along with the rest of the bookkeeping state. """ def __init__( self, block_count_total, max_kv_block_count, max_requests, block_size_tokens, max_seqlen ): super().__init__() - device = torch.cuda.current_device() - self.device = device + self.device = torch.cuda.current_device() self.max_blocks = block_count_total self.max_kv_blocks = max_kv_block_count self.max_bs = max_requests self.max_seqlen = max_seqlen - self._query_lengths_buf = torch.zeros(self.max_bs, dtype=torch.int32, device=device) - self._cu_query_seq_lengths_buf = torch.zeros( - self.max_bs + 1, dtype=torch.int32, device=device - ) - self._cu_kv_seq_lengths_buf = torch.zeros(self.max_bs + 1, dtype=torch.int32, device=device) - self._kv_seq_lengths_buf = torch.zeros(self.max_bs, dtype=torch.int32, device=device) - self._block_table_buf = torch.zeros( - (self.max_bs, self.max_kv_blocks), dtype=torch.int32, device=device - ) self._max_seqlen_q = 0 self._max_seqlen_k = 0 self.state_data = {} + # Set by bind_gpu_buffers(); references shared views in ContextGPUView._buf. + self._gpu_view = None - def update( - self, - request_query_lengths: torch.Tensor, - request_kv_length_offsets: torch.Tensor, - request_to_kv_block_ids: torch.Tensor, - batch_dimensions: InferenceBatchDimensions, - padded_batch_dimensions: InferenceBatchDimensions, - num_speculative_tokens: int = 0, - ): - """ - Args: - request_query_lengths: (>real_batch_size,) - request_kv_length_offsets: (>real_batch_size,) - request_to_kv_block_ids: (>real_batch_size, max_kv_blocks) - batch_dimensions: Configuration object containing real batch settings - padded_batch_dimensions: Configuration object containing padded batch settings - num_speculative_tokens: Number of speculative tokens - """ - # Extract values from configs - real_batch_size = batch_dimensions.req_count - padded_active_token_count = padded_batch_dimensions.token_count - padded_active_request_count = padded_batch_dimensions.req_count - - assert real_batch_size <= padded_active_request_count <= self.max_bs - assert request_query_lengths.shape[0] == real_batch_size - assert request_kv_length_offsets.shape[0] == real_batch_size - assert request_to_kv_block_ids.shape[0] == real_batch_size + def bind_gpu_buffers(self, gpu_view) -> None: + """Attach shared GPU buffer views from the context's ContextGPUView. - self.tensor_copy_and_pad( - self._query_lengths_buf, - request_query_lengths, - real_batch_size, - padded_active_request_count, - ) - self._cu_query_seq_lengths_buf[0] = 0 - self.tensor_copy_and_pad( - self._cu_query_seq_lengths_buf[1:], - torch.cumsum(request_query_lengths, dim=0), - real_batch_size, - padded_active_request_count, - is_cumulative_tensor=True, - ) - self.tensor_copy_and_pad( - self._kv_seq_lengths_buf, - request_kv_length_offsets + request_query_lengths, - real_batch_size, - padded_active_request_count, - ) - self.tensor_copy_and_pad( - self._block_table_buf, - request_to_kv_block_ids, - real_batch_size, - padded_active_request_count, - pad_value=torch.tensor(self.max_kv_blocks, dtype=torch.int32, device=self.device).fill_( - -1 - ), - ) - self._cu_kv_seq_lengths_buf[0] = 0 - self.tensor_copy_and_pad( - self._cu_kv_seq_lengths_buf[1:], - torch.cumsum(self._kv_seq_lengths_buf, dim=0), - real_batch_size, - padded_active_request_count, - is_cumulative_tensor=True, - ) - - if padded_batch_dimensions.prefill_req_count == 0: - self._max_seqlen_q = num_speculative_tokens + 1 - else: - # Make sure we will launch the prefill kernel for prefill graphs - self._max_seqlen_q = max(2, padded_batch_dimensions.token_count) + Called by :class:`DynamicInferenceContext` after ``self.gpu_view`` is + constructed. Both graphed and non-graphed MHA metadata bind to the + same views; only one is active per step, so sharing storage is safe. + """ + self._gpu_view = gpu_view - self._max_seqlen_k = self.max_seqlen + def set_state_data( + self, padded_active_request_count: int, max_seqlen_q: int, max_seqlen_k: int + ) -> None: + """Build ``state_data`` slices into the bound GPU buffers. + Called once per step from ``transfer_bookkeeping_to_gpu`` after the + coalesced H2D copy. No ``.copy_()`` calls, no kernel launches. + """ + assert self._gpu_view is not None, "bind_gpu_buffers() must be called first" + n = padded_active_request_count + v = self._gpu_view + self._max_seqlen_q = max_seqlen_q + self._max_seqlen_k = max_seqlen_k self.state_data = { - "query_lengths": self._query_lengths_buf[:padded_active_request_count], - "cu_query_seq_lengths": self._cu_query_seq_lengths_buf[ - : padded_active_request_count + 1 - ], - "cu_kv_seq_lengths": self._cu_kv_seq_lengths_buf[: padded_active_request_count + 1], - "kv_seq_lengths": self._kv_seq_lengths_buf[:padded_active_request_count], - "block_table": self._block_table_buf[0:padded_active_request_count, :], - "max_seqlen_q": self._max_seqlen_q, - "max_seqlen_k": self._max_seqlen_k, + "query_lengths": v.mha_query_lengths[:n], + "cu_query_seq_lengths": v.mha_cu_query_seq_lengths[: n + 1], + "cu_kv_seq_lengths": v.mha_cu_kv_seq_lengths[: n + 1], + "kv_seq_lengths": v.mha_kv_seq_lengths[:n], + "block_table": v.mha_block_table[:n, :], + "max_seqlen_q": max_seqlen_q, + "max_seqlen_k": max_seqlen_k, } def reset(self): + """Reset the metadata for the next batch. + + The GPU buffers live in the context's unified buffer and are fully + overwritten by the next H2D copy; clearing them here would launch + redundant CUDA kernels with no correctness benefit. """ - Reset the metadata for the next batch. - """ - self._query_lengths_buf.fill_(0) - self._cu_query_seq_lengths_buf.fill_(0) - self._cu_kv_seq_lengths_buf.fill_(0) - self._kv_seq_lengths_buf.fill_(0) - self._block_table_buf.fill_(0) self._max_seqlen_q = 0 self._max_seqlen_k = 0 class GraphedMHAMetadata(MHAMetadata): - """ - Metadata for MHA layer using flash-attention with CUDA graphs. - """ - - def __init__( - self, block_count_total, max_kv_block_count, max_requests, block_size_tokens, max_seqlen - ): - super().__init__( - block_count_total, max_kv_block_count, max_requests, block_size_tokens, max_seqlen - ) - - def update( - self, - request_query_lengths: torch.Tensor, - request_kv_length_offsets: torch.Tensor, - request_to_kv_block_ids: torch.Tensor, - batch_dimensions: InferenceBatchDimensions, - padded_batch_dimensions: InferenceBatchDimensions, - num_speculative_tokens: int = 0, - ): - """ - Args: - request_query_lengths: (>real_batch_size,) - request_kv_length_offsets: (>real_batch_size,) - request_to_kv_block_ids: (>real_batch_size, max_kv_blocks) - batch_dimensions: Configuration object containing real batch settings - padded_batch_dimensions: Configuration object containing padded batch settings - num_speculative_tokens: Number of speculative tokens - """ - super().update( - request_query_lengths, - request_kv_length_offsets, - request_to_kv_block_ids, - batch_dimensions, - padded_batch_dimensions, - num_speculative_tokens, - ) - - def reset(self): - super().reset() + """MHA metadata for CUDA-graphed execution.""" class NonGraphedMHAMetadata(MHAMetadata): - """ - Metadata for MHA layer using flash-attention without CUDA graphs. - """ - - def update( - self, - request_query_lengths: torch.Tensor, - request_kv_length_offsets: torch.Tensor, - request_to_kv_block_ids: torch.Tensor, - batch_dimensions: InferenceBatchDimensions, - padded_batch_dimensions: InferenceBatchDimensions, - num_speculative_tokens: int = 0, - ): - """ - Args: - request_query_lengths: (>real_batch_size,) - request_kv_length_offsets: (>real_batch_size,) - request_to_kv_block_ids: (>real_batch_size, max_kv_blocks) - batch_dimensions: Configuration object containing real batch settings - padded_batch_dimensions: Configuration object containing padded batch settings - num_speculative_tokens: Number of speculative tokens - """ - super().update( - request_query_lengths, - request_kv_length_offsets, - request_to_kv_block_ids, - batch_dimensions, - padded_batch_dimensions, - num_speculative_tokens, - ) - if len(self.state_data["query_lengths"]) > 0: - self.state_data["max_seqlen_q"] = torch.max(self.state_data["query_lengths"]).item() - self.state_data["max_seqlen_k"] = torch.max(self.state_data["kv_seq_lengths"]).item() - else: - self.state_data["max_seqlen_q"] = num_speculative_tokens + 1 - self.state_data["max_seqlen_k"] = 1 + """MHA metadata for non-graphed (eager) execution.""" diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 9765758eea7..159e1f90b34 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -46,6 +46,7 @@ from .attention_context.mamba_metadata import MambaMetadata from .attention_context.mha_metadata import GraphedMHAMetadata, NonGraphedMHAMetadata from .base_context import BaseInferenceContext +from .gpu_view import ContextGPUView from .kv_block_allocator import KVBlockAllocator from .mamba_slot_allocator import MambaSlotAllocator from .routing_metadata import RoutingMetadata @@ -327,6 +328,12 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC else: self.expert_model_parallel_group = None + # Optional CPU-side collective for EP batch-dimension sync. Populated by + # the engine via set_ep_zmq_communicator() when available. When set, + # match_graph_config() uses this to perform the MAX reduction on the + # CPU, avoiding a per-step NCCL AllReduce kernel on the compute stream. + self._ep_zmq_communicator = None + # Mamba states. mamba_inference_state_config = inference_config.mamba_inference_state_config self.is_hybrid_model = mamba_inference_state_config is not None @@ -635,6 +642,10 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC ep_group=self.expert_model_parallel_group, ) + self.smallest_non_decode_cuda_graph_size = min( + inference_config.cuda_graph_mixed_prefill_count, self.max_requests + ) + # Deal with chunked prefill self.enable_chunked_prefill = inference_config.enable_chunked_prefill @@ -764,8 +775,26 @@ def _allocate_mamba_states(self): self.mamba_metadata = MambaMetadata( max_requests=self.max_requests, max_tokens=self.max_tokens, + mamba_chunk_size=self.mamba_chunk_size, d_conv=self.mamba_conv_states_shape[-1], ) + # Bind the unified CPU/GPU buffers so the per-step Mamba metadata + # fields ride along with the single coalesced H2D in + # transfer_bookkeeping_to_gpu(). + self.mamba_metadata.bind_cpu_buffers( + { + "batch_indices_decode": self._cpu_mamba_batch_indices_decode, + "batch_indices_prefill": self._cpu_mamba_batch_indices_prefill, + "seq_idx": self._cpu_mamba_seq_idx, + "cu_seqlens": self._cpu_mamba_cu_seqlens, + "cu_chunk_seqlens": self._cpu_mamba_cu_chunk_seqlens, + "last_chunk_indices": self._cpu_mamba_last_chunk_indices, + "seq_idx_for_varlen": self._cpu_mamba_seq_idx_for_varlen, + "conv_seq_idx": self._cpu_mamba_conv_seq_idx, + "conv_seq_start": self._cpu_mamba_conv_seq_start, + } + ) + self.mamba_metadata.bind_gpu_buffers(self.gpu_view) self.mamba_conv_states = torch.empty( (self.num_mamba_layers, self.max_requests) + self.mamba_conv_states_shape, dtype=self.mamba_conv_states_dtype, @@ -841,54 +870,191 @@ def initialize_all_tensors(self) -> None: f"Please move tensor '{key}'." ) - # Per-request state. + # Per-request state (CPU, pinned memory for fast H2D transfer). self.request_ids = torch.full( - (self.max_requests,), -1, dtype=torch.int32, device=torch.cuda.current_device() + (self.max_requests,), -1, dtype=torch.int32, device='cpu', pin_memory=True ) # request_query_lengths is the input prompt tokens length during prefill phase (1st step) and then 1 for the decode phase (i.e During generation) - self.request_query_lengths = torch.empty_like(self.request_ids) + self.request_query_lengths = torch.empty( + self.max_requests, dtype=torch.int32, device='cpu', pin_memory=True + ) # True only for a new request , then after a forward pass it is set to False - self.request_in_prefill_status_tensor = torch.empty_like(self.request_ids) + self.request_in_prefill_status_tensor = torch.empty( + self.max_requests, dtype=torch.int32, device='cpu', pin_memory=True + ) # request_output_lengths is len(input_prompt_tokens) + num_tokens_to_generate - self.request_output_lengths = torch.empty_like(self.request_ids) + self.request_output_lengths = torch.empty( + self.max_requests, dtype=torch.int32, device='cpu', pin_memory=True + ) # request_kv_length_offsets is the same as query length during prefill phase (1st step) and then 1 for the decode phase (i.e During generation) - self.request_kv_length_offsets = torch.empty_like(self.request_ids) - self.request_kv_block_counts = torch.empty_like(self.request_ids) - self.request_last_kv_block_id = torch.empty_like(self.request_ids) + self.request_kv_length_offsets = torch.empty( + self.max_requests, dtype=torch.int32, device='cpu', pin_memory=True + ) + self.request_kv_block_counts = torch.empty( + self.max_requests, dtype=torch.int32, device='cpu', pin_memory=True + ) + self.request_last_kv_block_id = torch.empty( + self.max_requests, dtype=torch.int32, device='cpu', pin_memory=True + ) # request_last_kv_block_offset represents number of tokens in the last kv block - self.request_last_kv_block_offset = torch.empty_like(self.request_ids) + self.request_last_kv_block_offset = torch.empty( + self.max_requests, dtype=torch.int32, device='cpu', pin_memory=True + ) self.request_to_kv_block_ids = torch.full( (self.max_requests, self.max_kv_block_count), -1, dtype=torch.int, - device=torch.cuda.current_device(), + device='cpu', + pin_memory=True, ) - # Track request metadata. + # Track request metadata. Backed by pinned CPU memory: bookkeeping is + # CPU-resident; GPU consumers read from the active-slice mirror in + # `active_request_metadata` (also CPU pinned, refreshed each step). self.request_metadata = { - label: torch.empty( - (self.max_requests,), dtype=dtype, device=torch.cuda.current_device() - ) + label: torch.empty((self.max_requests,), dtype=dtype, device='cpu', pin_memory=True) for label, dtype in self.request_metadata_types } - # Per-token state. - self.token_to_input_ids = torch.full( - (self.max_tokens,), 0, dtype=torch.long, device=torch.cuda.current_device() - ) - self.token_to_pos_ids = torch.full_like(self.token_to_input_ids, 0) - self.token_to_request_idx = torch.empty_like(self.token_to_input_ids) - self.token_to_block_idx = torch.empty_like(self.token_to_input_ids) + # Static tensor addresses of active slices to enable fast inference + # kernels. Pinned CPU mirrors of `request_metadata`, refreshed each + # step by `build_active_slices()` from the active subrange. + self.active_request_metadata = { + label: torch.empty_like(tensor, pin_memory=True) + for label, tensor in self.request_metadata.items() + } + + # Coalesced pinned CPU buffer for the bookkeeping fields that get + # transferred to GPU each step via transfer_bookkeeping_to_gpu(). + # Layout matches ContextGPUView._buf so a single cudaMemcpyAsync + # suffices. Int64 token fields come first (8-byte aligned automatically), + # then int32 token fields, then int32 request-staging fields. + # token_to_input_ids (int64, max_tokens) + # token_to_pos_ids (int64, max_tokens) + # token_to_block_idx (int32, max_tokens) + # token_to_local_position_within_kv_block (int32, max_tokens) + # token_to_request_idx (int32, max_tokens) + # token_to_position_in_request (int32, max_tokens) + # request_in_prefill_status (staging) (int32, max_requests) + # request_query_lengths (staging) (int32, max_requests) + # request_kv_length_offsets (staging) (int32, max_requests) + # + # Token fields are aliased with the source-of-truth attributes + # (`self.token_to_input_ids`, etc.) because the forward pass reads + # `gpu_view.token_to_input_ids[:n_tok]` which matches the CPU slot + # layout `[0, n_tok)`. Request fields, however, are read on GPU at + # `[:n_active]` but on CPU at `[paused_count:total_count)` — so the + # staging slots here are refreshed each step by copying the active + # slice from the persistent `request_*` tensors above. + _tok_int64_bytes = self.max_tokens * 8 + _tok_int32_bytes = self.max_tokens * 4 + _req_int32_bytes = self.max_requests * 4 + # MHA section: 5 fields (int32) shared between GraphedMHAMetadata and + # NonGraphedMHAMetadata. max_bs == max_requests. + _mha_query_lengths_bytes = self.max_requests * 4 + _mha_cu_query_seq_lengths_bytes = (self.max_requests + 1) * 4 + _mha_kv_seq_lengths_bytes = self.max_requests * 4 + _mha_cu_kv_seq_lengths_bytes = (self.max_requests + 1) * 4 + _mha_block_table_bytes = self.max_requests * self.max_kv_block_count * 4 + # Mamba section: 9 int32 fields (hybrid models only). Must match the + # MambaMetadata shapes (mirrors the layout documented in ContextGPUView). + if self.is_hybrid_model: + self._max_mamba_chunks = self.max_tokens // self.mamba_chunk_size + self.max_requests + _mamba_batch_indices_decode_bytes = self.max_requests * 4 + _mamba_batch_indices_prefill_bytes = self.max_requests * 4 + _mamba_seq_idx_bytes = self.max_tokens * 4 + _mamba_cu_seqlens_bytes = (self.max_requests + 1) * 4 + _mamba_cu_chunk_seqlens_bytes = (self._max_mamba_chunks + 1) * 4 + _mamba_last_chunk_indices_bytes = self.max_requests * 4 + _mamba_seq_idx_for_varlen_bytes = self._max_mamba_chunks * 4 + _mamba_conv_seq_idx_bytes = self.max_tokens * 4 + _mamba_conv_seq_start_bytes = self.max_tokens * 4 + else: + self._max_mamba_chunks = 0 + _mamba_batch_indices_decode_bytes = 0 + _mamba_batch_indices_prefill_bytes = 0 + _mamba_seq_idx_bytes = 0 + _mamba_cu_seqlens_bytes = 0 + _mamba_cu_chunk_seqlens_bytes = 0 + _mamba_last_chunk_indices_bytes = 0 + _mamba_seq_idx_for_varlen_bytes = 0 + _mamba_conv_seq_idx_bytes = 0 + _mamba_conv_seq_start_bytes = 0 + _total_bytes = ( + 2 * _tok_int64_bytes + + 4 * _tok_int32_bytes + + 3 * _req_int32_bytes + + _mha_query_lengths_bytes + + _mha_cu_query_seq_lengths_bytes + + _mha_kv_seq_lengths_bytes + + _mha_cu_kv_seq_lengths_bytes + + _mha_block_table_bytes + + _mamba_batch_indices_decode_bytes + + _mamba_batch_indices_prefill_bytes + + _mamba_seq_idx_bytes + + _mamba_cu_seqlens_bytes + + _mamba_cu_chunk_seqlens_bytes + + _mamba_last_chunk_indices_bytes + + _mamba_seq_idx_for_varlen_bytes + + _mamba_conv_seq_idx_bytes + + _mamba_conv_seq_start_bytes + ) + self._cpu_bookkeeping_buf = torch.empty( + _total_bytes, dtype=torch.uint8, device='cpu', pin_memory=True + ) + # token_to_input_ids and token_to_pos_ids were previously torch.full(0); + # zero the whole buffer so their views start at 0 too, and so the + # request staging slots start with a deterministic value. + self._cpu_bookkeeping_buf.fill_(0) + + _off = 0 + # Per-token state (source-of-truth lives in the coalesced buffer since + # the CPU-side bookkeeping and the GPU forward pass use the same + # `[:n_tok]` slice). + self.token_to_input_ids = self._cpu_bookkeeping_buf[_off : _off + _tok_int64_bytes].view( + torch.long + ) + _off += _tok_int64_bytes + self.token_to_pos_ids = self._cpu_bookkeeping_buf[_off : _off + _tok_int64_bytes].view( + torch.long + ) + _off += _tok_int64_bytes + self.token_to_block_idx = self._cpu_bookkeeping_buf[_off : _off + _tok_int32_bytes].view( + torch.int32 + ) + _off += _tok_int32_bytes # i.e For a set of tokens A B C D E F .. and block_size 4: # token_to_position_in_request is [0, 1, 2, 3, 4, 5] # token_to_local_position_within_kv_block is [0 , 1, 2, 3, 0, 1, 2] - self.token_to_position_in_request = torch.empty_like(self.token_to_input_ids) - self.token_to_local_position_within_kv_block = torch.empty_like(self.token_to_input_ids) + self.token_to_local_position_within_kv_block = self._cpu_bookkeeping_buf[ + _off : _off + _tok_int32_bytes + ].view(torch.int32) + _off += _tok_int32_bytes + self.token_to_request_idx = self._cpu_bookkeeping_buf[_off : _off + _tok_int32_bytes].view( + torch.int32 + ) + _off += _tok_int32_bytes + self.token_to_position_in_request = self._cpu_bookkeeping_buf[ + _off : _off + _tok_int32_bytes + ].view(torch.int32) + _off += _tok_int32_bytes + + # Request-level staging views into the coalesced buffer. Write-only on + # CPU (refreshed from persistent tensors in transfer_bookkeeping_to_gpu); + # read-only on GPU via matching slots in ContextGPUView._buf. + self._staging_request_in_prefill_status = self._cpu_bookkeeping_buf[ + _off : _off + _req_int32_bytes + ].view(torch.int32) + _off += _req_int32_bytes + self._staging_request_query_lengths = self._cpu_bookkeeping_buf[ + _off : _off + _req_int32_bytes + ].view(torch.int32) + _off += _req_int32_bytes + self._staging_request_kv_length_offsets = self._cpu_bookkeeping_buf[ + _off : _off + _req_int32_bytes + ].view(torch.int32) + _off += _req_int32_bytes - # Static tensor addresses of active slices to enable fast inference kernels. - self.active_request_metadata = { - label: torch.empty_like(tensor) for label, tensor in self.request_metadata.items() - } # Static tensor addresses to make `last_token_logits` graphable with speculative decoding. max_logit_idxs = self.max_requests * (self.num_speculative_tokens + 1) self.active_logit_idxs = torch.zeros( @@ -898,14 +1064,97 @@ def initialize_all_tensors(self) -> None: max_logit_idxs, dtype=torch.int32, device=torch.cuda.current_device() ) - # NOTE: Need to build this outside the UVM / TMS context to avoid IMA. + # MHA flash-attention metadata views (write-only on CPU, read-only on + # GPU via the matching region of ContextGPUView._buf). Populated per + # step by initialize_attention_state(); transferred as part of the + # single coalesced H2D in transfer_bookkeeping_to_gpu(). + self._cpu_mha_query_lengths = self._cpu_bookkeeping_buf[ + _off : _off + _mha_query_lengths_bytes + ].view(torch.int32) + _off += _mha_query_lengths_bytes + self._cpu_mha_cu_query_seq_lengths = self._cpu_bookkeeping_buf[ + _off : _off + _mha_cu_query_seq_lengths_bytes + ].view(torch.int32) + _off += _mha_cu_query_seq_lengths_bytes + self._cpu_mha_kv_seq_lengths = self._cpu_bookkeeping_buf[ + _off : _off + _mha_kv_seq_lengths_bytes + ].view(torch.int32) + _off += _mha_kv_seq_lengths_bytes + self._cpu_mha_cu_kv_seq_lengths = self._cpu_bookkeeping_buf[ + _off : _off + _mha_cu_kv_seq_lengths_bytes + ].view(torch.int32) + _off += _mha_cu_kv_seq_lengths_bytes + self._cpu_mha_block_table = ( + self._cpu_bookkeeping_buf[_off : _off + _mha_block_table_bytes] + .view(torch.int32) + .view(self.max_requests, self.max_kv_block_count) + ) + _off += _mha_block_table_bytes + + # Mamba varlen metadata views (hybrid models only). Populated per step + # by MambaMetadata.compute_cpu_metadata(); transferred as part of the + # single coalesced H2D in transfer_bookkeeping_to_gpu(). if self.is_hybrid_model: - self.mamba_metadata = MambaMetadata( - max_requests=self.max_requests, - max_tokens=self.max_tokens, - mamba_chunk_size=self.mamba_chunk_size, - d_conv=self.mamba_conv_states_shape[-1], + self._cpu_mamba_batch_indices_decode = self._cpu_bookkeeping_buf[ + _off : _off + _mamba_batch_indices_decode_bytes + ].view(torch.int32) + _off += _mamba_batch_indices_decode_bytes + self._cpu_mamba_batch_indices_prefill = self._cpu_bookkeeping_buf[ + _off : _off + _mamba_batch_indices_prefill_bytes + ].view(torch.int32) + _off += _mamba_batch_indices_prefill_bytes + self._cpu_mamba_seq_idx = ( + self._cpu_bookkeeping_buf[_off : _off + _mamba_seq_idx_bytes] + .view(torch.int32) + .view(1, self.max_tokens) ) + _off += _mamba_seq_idx_bytes + self._cpu_mamba_cu_seqlens = self._cpu_bookkeeping_buf[ + _off : _off + _mamba_cu_seqlens_bytes + ].view(torch.int32) + _off += _mamba_cu_seqlens_bytes + self._cpu_mamba_cu_chunk_seqlens = self._cpu_bookkeeping_buf[ + _off : _off + _mamba_cu_chunk_seqlens_bytes + ].view(torch.int32) + _off += _mamba_cu_chunk_seqlens_bytes + self._cpu_mamba_last_chunk_indices = self._cpu_bookkeeping_buf[ + _off : _off + _mamba_last_chunk_indices_bytes + ].view(torch.int32) + _off += _mamba_last_chunk_indices_bytes + self._cpu_mamba_seq_idx_for_varlen = self._cpu_bookkeeping_buf[ + _off : _off + _mamba_seq_idx_for_varlen_bytes + ].view(torch.int32) + _off += _mamba_seq_idx_for_varlen_bytes + self._cpu_mamba_conv_seq_idx = self._cpu_bookkeeping_buf[ + _off : _off + _mamba_conv_seq_idx_bytes + ].view(torch.int32) + _off += _mamba_conv_seq_idx_bytes + self._cpu_mamba_conv_seq_start = self._cpu_bookkeeping_buf[ + _off : _off + _mamba_conv_seq_start_bytes + ].view(torch.int32) + _off += _mamba_conv_seq_start_bytes + + assert _off == _total_bytes, f"layout bug: wrote {_off} of {_total_bytes} bytes" + + # GPU view: the single interface for GPU code to read context state. + # Populated per-step by transfer_bookkeeping_to_gpu(). + self.gpu_view = ContextGPUView( + max_requests=self.max_requests, + max_tokens=self.max_tokens, + max_kv_blocks=self.max_kv_block_count, + device=torch.cuda.current_device(), + max_mamba_chunks=self._max_mamba_chunks, + ) + + # Bind the shared MHA GPU views to both graph and non-graph metadata; + # only one is active per step, so sharing storage is safe. + self.graph_attn_metadata["mha_metadata"].bind_gpu_buffers(self.gpu_view) + self.non_graph_attn_metadata["mha_metadata"].bind_gpu_buffers(self.gpu_view) + + # Deferred Mamba GPU operations. Populated by add_request() / + # update_requests() (CPU phase), executed by transfer_bookkeeping_to_gpu(). + self._pending_mamba_zeros: list = [] + self._pending_mamba_restores: list = [] # Allocate large non-graphed buffers. need_static_addr = ( @@ -1131,12 +1380,11 @@ def pad_active_slices(self): active_decode_token_count : active_decode_token_count + active_prefill_count ] prefill_idxs = self.paused_request_count + active_decode_count - torch.cumsum( - self.request_query_lengths[prefill_idxs : self.total_request_count], - dim=0, - out=prefill_dst, - ) - prefill_dst.add_(active_decode_token_count - 1) + prefill_lengths = self.request_query_lengths[prefill_idxs : self.total_request_count] + if active_prefill_count > 0: + prefill_cumsum = torch.cumsum(prefill_lengths, dim=0, dtype=torch.int32) + prefill_cumsum.add_(active_decode_token_count - 1) + prefill_dst.copy_(prefill_cumsum, non_blocking=True) self.active_logit_idxs[active_decode_token_count + active_prefill_count :].zero_() @@ -1158,12 +1406,12 @@ def append_key_value_cache(self, layer_number: int, key: Tensor, value: Tensor) value=value, memory_buffer=self.memory_buffer, padded_active_token_count=self.padded_active_token_count, - token_to_block_idx=self.token_to_block_idx, - token_to_local_position_within_kv_block=self.token_to_local_position_within_kv_block, + token_to_block_idx=self.gpu_view.token_to_block_idx, + token_to_local_position_within_kv_block=self.gpu_view.token_to_local_position_within_kv_block, ) - block_idx = self.token_to_block_idx[: self.padded_active_token_count] - local_kv_seq_idx = self.token_to_local_position_within_kv_block[ + block_idx = self.gpu_view.token_to_block_idx[: self.padded_active_token_count] + local_kv_seq_idx = self.gpu_view.token_to_local_position_within_kv_block[ : self.padded_active_token_count ] @@ -1299,7 +1547,7 @@ def apply_fused_qk_rotary_emb( # use .view instead of .reshape to avoid extra transpose operations query_rope, key_rope = flashinfer.rope.apply_rope_with_cos_sin_cache( - positions=self.token_to_pos_ids[:n], + positions=self.gpu_view.token_to_pos_ids[:n], query=query[:n].reshape(n, num_q_heads * head_size), key=key[:n].reshape(n, num_k_heads * head_size), head_size=head_size, @@ -1332,7 +1580,7 @@ def apply_rotary_emb_query( (Tensor) Query tensor after applying rotary embeddings. """ n = self.padded_active_token_count - query_seq_idx = self.token_to_pos_ids[:n] + query_seq_idx = self.gpu_view.token_to_pos_ids[:n] query_emb = query_emb[query_seq_idx] query[:n] = apply_rotary_pos_emb( t=query[:n], @@ -1364,7 +1612,7 @@ def apply_rotary_emb_key( (Tensor) Key tensor after applying rotary embeddings. """ n = self.padded_active_token_count - key_seq_idx = self.token_to_position_in_request[:n] + key_seq_idx = self.gpu_view.token_to_position_in_request[:n] key_emb = key_emb[key_seq_idx] if self.is_decode_only(): if key.shape[0] != n: @@ -1383,6 +1631,20 @@ def apply_rotary_emb_key( ) return key + def set_ep_zmq_communicator(self, communicator) -> None: + """Attach an EP-group ZMQ communicator for CPU-side sync collectives. + + When set, match_graph_config() uses this communicator's + sync_all_reduce_max() to perform the EP batch-dimension MAX reduction on + the CPU instead of launching a NCCL AllReduce kernel on the compute + stream. Expected to be called once by the inference engine after both + the context and the communicator have been created. + + Args: + communicator: AsyncZMQCommunicator over the EP process group. + """ + self._ep_zmq_communicator = communicator + def reset_attention_state(self) -> None: """Reset state used within attention, after each step.""" # Attention metadata reset is now handled by MHAMetadata.reset() @@ -1471,7 +1733,7 @@ def add_dummy_requests_parallel( self.request_kv_block_counts[request_slice] = block_counts for i, (label, dtype) in enumerate(self.request_metadata_types): self.request_metadata[label][request_slice] = torch.tensor( - metadata_cols[i], dtype=dtype, device=torch.cuda.current_device() + metadata_cols[i], dtype=dtype, device='cpu' ) dummy_block_idx = self.kv_block_allocator.dummy_block_idx @@ -1530,8 +1792,7 @@ def add_dummy_requests_parallel( raise ContextOverflowError( requests[logical_idx].request_id, "No Mamba slots available" ) - self.mamba_conv_states[:, mamba_idx] = 0.0 - self.mamba_ssm_states[:, mamba_idx] = 0.0 + self._pending_mamba_zeros.append(mamba_idx) self.mamba_metadata.request_to_mamba_state_idx[request_idx] = mamba_idx self.active_token_count = token_end @@ -1553,7 +1814,7 @@ def add_dummy_requests_for_cudagraph_capture( # Pre-construct shared objects (safe due to deep copy in DynamicInferenceRequest.__post_init__) shared_sampling_params = SamplingParams(num_tokens_to_generate=1, termination_id=-1) shared_decode_tokens = torch.zeros( - self.num_speculative_tokens + 1, dtype=torch.long, device=torch.cuda.current_device() + self.num_speculative_tokens + 1, dtype=torch.long, device='cpu' ) decode_requests = [ @@ -1583,9 +1844,7 @@ def add_dummy_requests_for_cudagraph_capture( assert per_prefill_tokens > 0 # Create a single large tensor and slice from it for each prefill request max_prefill_tokens = per_prefill_tokens + (1 if rem_prefill_tokens > 0 else 0) - shared_prefill_tokens = torch.zeros( - max_prefill_tokens, dtype=torch.long, device=torch.cuda.current_device() - ) + shared_prefill_tokens = torch.zeros(max_prefill_tokens, dtype=torch.long, device='cpu') prefill_requests = [ DynamicInferenceRequest( @@ -1631,7 +1890,7 @@ def add_dummy_requests_for_expert_parallel_step( self.active_token_count = T self.num_prefill_requests = N_prefill - # 2. Per-request state consumed by mha_metadata.update(). + # 2. Per-request state consumed by initialize_attention_state(). # Decode requests come first, followed by prefill requests. self.request_query_lengths[0:N_decode].fill_(tokens_per_decode_request) if N_prefill > 0: @@ -1690,6 +1949,10 @@ def initialize_attention_state( Return: None. """ + # Launch deferred Mamba GPU ops first (state zeroing/restore) so they + # overlap with the CPU work below. These are non-blocking GPU kernels. + self._execute_pending_mamba_ops() + self.is_creating_cuda_graphs = construct_graph_dimensions is not None assert not ( self.is_creating_cuda_graphs and is_expert_parallel_dummy_cuda_graph_step @@ -1719,8 +1982,12 @@ def initialize_attention_state( best_graph = CUDAGraphBatchDimensionBuilder.match_graph_config( batch_dimensions, self.cuda_graph_batch_dimensions_list, + smallest_non_decode_cuda_graph_size=self.smallest_non_decode_cuda_graph_size, strict=self.is_hybrid_model, + decode_only_cuda_graphs=(not self.use_cuda_graphs_for_non_decode_steps), ep_group=self.expert_model_parallel_group, + num_speculative_tokens=self.num_speculative_tokens, + ep_zmq_communicator=self._ep_zmq_communicator, match_ep_token_counts=self._nccl_ep_dispatcher, ) self._using_cuda_graph_this_step = best_graph is not None @@ -1804,31 +2071,98 @@ def initialize_attention_state( ) assert self.active_attn_metadata is not None - self.active_attn_metadata["mha_metadata"].update( - request_query_lengths=query_lengths_view, - request_kv_length_offsets=request_kv_length_offsets_view, - request_to_kv_block_ids=request_to_kv_block_ids_view, - batch_dimensions=attn_dimensions, - padded_batch_dimensions=self.padded_batch_dimensions, - num_speculative_tokens=self.num_speculative_tokens, + + # Compute MHA metadata directly into the pinned CPU section of + # _cpu_bookkeeping_buf. The single coalesced H2D in + # transfer_bookkeeping_to_gpu() covers these fields along with the rest + # of the bookkeeping state, so no ephemeral tensors and no per-field + # cudaMemcpyAsyncs. + real_bs = attn_dimensions.req_count + padded_bs = self.padded_batch_dimensions.req_count + mha = self.active_attn_metadata["mha_metadata"] + + # Query lengths: [0:real_bs] real data, [real_bs:padded_bs] zero pad. + self._cpu_mha_query_lengths[:real_bs] = query_lengths_view[:real_bs] + if real_bs < padded_bs: + self._cpu_mha_query_lengths[real_bs:padded_bs] = 0 + + # Cumulative query lengths (padded slots repeat cu[real_bs]). + self._cpu_mha_cu_query_seq_lengths[0] = 0 + if real_bs > 0: + self._cpu_mha_cu_query_seq_lengths[1 : real_bs + 1] = torch.cumsum( + query_lengths_view[:real_bs], dim=0 + ) + if real_bs < padded_bs: + self._cpu_mha_cu_query_seq_lengths[real_bs + 1 : padded_bs + 1] = ( + self._cpu_mha_cu_query_seq_lengths[real_bs] + ) + + # KV sequence lengths: [0:real_bs] = kv_offsets + query_lengths. + self._cpu_mha_kv_seq_lengths[:real_bs] = ( + request_kv_length_offsets_view[:real_bs] + query_lengths_view[:real_bs] + ) + if real_bs < padded_bs: + self._cpu_mha_kv_seq_lengths[real_bs:padded_bs] = 0 + + # Cumulative KV lengths. + self._cpu_mha_cu_kv_seq_lengths[0] = 0 + if real_bs > 0: + self._cpu_mha_cu_kv_seq_lengths[1 : real_bs + 1] = torch.cumsum( + self._cpu_mha_kv_seq_lengths[:real_bs], dim=0 + ) + if real_bs < padded_bs: + self._cpu_mha_cu_kv_seq_lengths[real_bs + 1 : padded_bs + 1] = ( + self._cpu_mha_cu_kv_seq_lengths[real_bs] + ) + + # Block table: [0:real_bs] real, [real_bs:padded_bs] = -1 sentinel. + self._cpu_mha_block_table[:real_bs] = request_to_kv_block_ids_view[:real_bs] + if real_bs < padded_bs: + self._cpu_mha_block_table[real_bs:padded_bs] = -1 + + # Max sequence lengths (Python scalars; consumed as kernel launch args). + if not self.using_cuda_graph_this_step() and real_bs > 0: + # NonGraphedMHAMetadata: use actual max values. + max_seqlen_q = self._cpu_mha_query_lengths[:real_bs].max().item() + max_seqlen_k = self._cpu_mha_kv_seq_lengths[:real_bs].max().item() + else: + # GraphedMHAMetadata: use conservative bounds. + if self.padded_batch_dimensions.prefill_req_count == 0: + max_seqlen_q = self.num_speculative_tokens + 1 + else: + max_seqlen_q = max(2, self.padded_batch_dimensions.token_count) + max_seqlen_k = mha.max_seqlen + if not self.using_cuda_graph_this_step() and real_bs == 0: + max_seqlen_q = self.num_speculative_tokens + 1 + max_seqlen_k = 1 + + # Bind state_data to GPU views now. set_state_data() only creates Python + # slice references into the GPU buffer (no GPU reads), so it's safe to + # call before the H2D in transfer_bookkeeping_to_gpu(). This guarantees + # that callers reading state_data["block_table"] etc. between + # initialize_attention_state() and transfer_bookkeeping_to_gpu() see + # populated entries (the actual data fill happens at the H2D). + mha.set_state_data( + padded_active_request_count=padded_bs, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, ) if self.is_hybrid_model: - active_mamba_indices_view = self.mamba_metadata.request_to_mamba_state_idx[active_slice] - token_to_request_idx_view = self.token_to_request_idx[: self.active_token_count] - cu_seqlens = self.active_attn_metadata["mha_metadata"].state_data[ - "cu_query_seq_lengths" - ] + # Mamba metadata update is deferred to transfer_bookkeeping_to_gpu() + # because it writes to GPU buffers. Store the parameters here. + # intermediate_offsets_gpu / intermediate_counts_gpu get the CPU-side + # slices here; H2D transfer happens in transfer_bookkeeping_to_gpu(). intermediate_offsets_gpu = None intermediate_counts_gpu = None if self.mamba_slot_allocator is not None: intermediate_offsets_gpu, intermediate_counts_gpu = ( - self.mamba_slot_allocator.get_intermediate_gpu_data() + self.mamba_slot_allocator.get_intermediate_cpu_data() ) - self.mamba_metadata.update( - active_mamba_indices_view, - token_to_request_idx_view, - cu_seqlens, + self._pending_mamba_transfer = self.mamba_metadata.compute_cpu_metadata( + active_mamba_indices=self.mamba_metadata.request_to_mamba_state_idx[active_slice], + token_to_request_idx=self.token_to_request_idx[: self.active_token_count], + cpu_cu_query=self._cpu_mha_cu_query_seq_lengths, batch_dimensions=attn_dimensions, padded_batch_dimensions=self.padded_batch_dimensions, enable_chunked_prefill=self.is_chunked_prefill_enabled(), @@ -1847,8 +2181,97 @@ def initialize_attention_state( if self._nccl_ep_dispatcher: NCCLAllGatherDispatcher._use_allgather_v = not self.using_cuda_graph_this_step() + # Flush any Mamba ops queued by add_dummy_requests_for_cudagraph_capture + # (warmup) or add_dummy_requests_for_expert_parallel_step (EP dummy step). + # The earlier call at the top drained ops queued by add_request() before + # this function ran; this call covers ops queued during the function. + # No-op when the queue is already empty (regular non-warmup steps). + self._execute_pending_mamba_ops() + + # Run the H2D transfer here so callers that bypass the controller + # (e.g. unit tests that call `model.forward()` directly after + # `initialize_attention_state()`) see populated GPU bookkeeping. The + # text-generation controller still calls `transfer_bookkeeping_to_gpu` + # explicitly; that second call is a cheap idempotent re-copy. + self.transfer_bookkeeping_to_gpu() + + def _execute_pending_mamba_ops(self) -> None: + """Execute Mamba GPU operations deferred from add_request() / update_requests(). + + This runs at the start of initialize_attention_state() so that all GPU + Mamba state is correct before the forward pass. + """ + if not (self._pending_mamba_restores or self._pending_mamba_zeros): + return + + # Restore cached Mamba state to live buffers. On failure, fall back to zeroing. + for request_idx, block_id, mamba_idx in self._pending_mamba_restores: + restored = self.mamba_slot_allocator.restore_to_live(request_idx, block_id) + if not restored: + self._pending_mamba_zeros.append(mamba_idx) + self._pending_mamba_restores.clear() + + # Batch-zero newly allocated Mamba slots. + if self._pending_mamba_zeros: + device = self.mamba_conv_states.device + indices = torch.tensor(self._pending_mamba_zeros, dtype=torch.long, device=device) + self.mamba_conv_states[:, indices] = 0.0 + self.mamba_ssm_states[:, indices] = 0.0 + self._pending_mamba_zeros.clear() + + def transfer_bookkeeping_to_gpu(self) -> None: + """Batch transfer CPU bookkeeping state to GPU staging buffers. + + Called after initialize_attention_state() and before the forward pass. + All copies use non_blocking=True with pinned CPU memory. CUDA stream + ordering guarantees the forward pass sees completed transfers. + + The 9 bookkeeping fields are backed by one contiguous pinned CPU buffer + and one contiguous GPU buffer; a single cudaMemcpyAsync suffices. + Request-level staging slots are refreshed from the persistent CPU + tensors immediately before the H2D (GPU reads them at `[:n_active]` + while CPU bookkeeping keeps them at `[paused_count:total_count)`). + """ + n_active = self.total_request_count - self.paused_request_count + active_slice = slice(self.paused_request_count, self.total_request_count) + padded_active = max(n_active, self.padded_active_request_count) + + # Refresh request-level staging slots from the persistent CPU source. + # CPU-to-CPU slice assignment on pinned memory (~7.5 KB total for 3 + # int32 fields at max_requests=624). Negligible vs. the launch overhead + # we save by merging 9 H2D memcpys into 1. + self._staging_request_in_prefill_status[:n_active] = self.request_in_prefill_status_tensor[ + active_slice + ] + self._staging_request_query_lengths[:n_active] = self.request_query_lengths[active_slice] + self._staging_request_kv_length_offsets[:n_active] = self.request_kv_length_offsets[ + active_slice + ] + # Full-iteration CUDA graphs may have captured GPU consumers with the + # padded graph request count. Keep those padded staging rows bounded so + # graph replay never builds indices from stale request lengths. + if n_active < padded_active: + self._staging_request_in_prefill_status[n_active:padded_active] = 0 + self._staging_request_query_lengths[n_active:padded_active] = 0 + self._staging_request_kv_length_offsets[n_active:padded_active] = 0 + + # Coalesced H2D: one cudaMemcpyAsync for the entire bookkeeping buffer. + # Copying the whole (max_tokens + max_requests)-sized buffer including + # unused slots is cheap (~71 KB total, ~3-5 us on PCIe Gen4) and saves + # 8 redundant launch overheads vs. the prior per-field copies. + self.gpu_view._buf.copy_(self._cpu_bookkeeping_buf, non_blocking=True) + + # MHA metadata GPU views were already bound to state_data in + # initialize_attention_state(); the H2D above populates the underlying + # bytes. Nothing else to do here for MHA. + + # Mamba metadata: copy pre-computed CPU tensors to GPU buffers. + if hasattr(self, '_pending_mamba_transfer') and self._pending_mamba_transfer is not None: + self.mamba_metadata.load_from_cpu(self._pending_mamba_transfer) + self._pending_mamba_transfer = None + def reset_tensors(self) -> None: - """Fill all GPU tensors with sentinel values.""" + """Fill all bookkeeping tensors with sentinel values.""" # Reset request indexes. self.request_ids.fill_(-1) @@ -1952,8 +2375,8 @@ def current_input_and_position_ids( self.num_speculative_tokens + 1 ) return ( - self.token_to_input_ids[:num_tokens].unsqueeze(0), - self.token_to_pos_ids[:num_tokens].unsqueeze(0), + self.gpu_view.token_to_input_ids[:num_tokens].unsqueeze(0), + self.gpu_view.token_to_pos_ids[:num_tokens].unsqueeze(0), ) def speculative_required_logit_indices(self) -> Tensor: @@ -2228,9 +2651,7 @@ def add_request( # Increment ref counts and update timestamps for matched (shared) blocks if num_matched_blocks > 0: - matched_tensor = torch.tensor( - matched_block_ids, dtype=torch.int32, device=torch.cuda.current_device() - ) + matched_tensor = torch.tensor(matched_block_ids, dtype=torch.int32, device='cpu') self.kv_block_allocator.block_ref_counts[matched_tensor] += 1 if self.prefix_caching_eviction_policy == PrefixCachingEvictionPolicy.LRU: self.kv_block_allocator.update_timestamps(matched_tensor) @@ -2354,17 +2775,18 @@ def _register_range(start: int, end: int): # Restore Mamba state from the block corresponding to prefix_skip_tokens restore_block_count = prefix_skip_tokens // self.block_size_tokens - restored = False if restore_block_count > 0 and self.mamba_slot_allocator is not None: restore_block_id = matched_block_ids[restore_block_count - 1] - restored = self.mamba_slot_allocator.restore_to_live( - self.total_request_count, restore_block_id + self._pending_mamba_restores.append( + (self.total_request_count, restore_block_id, mamba_idx) ) - if not restored: - self.mamba_conv_states[:, mamba_idx] = 0.0 - self.mamba_ssm_states[:, mamba_idx] = 0.0 + else: + self._pending_mamba_zeros.append(mamba_idx) - # Compute intermediate offsets for state extraction during forward pass + # compute_and_store_offsets sets both CPU state (hash_to_block_id, + # _eos_cache_block_id_gpu) and GPU staging buffers. Runs immediately + # because commit_intermediate_states() reads the CPU state after the + # forward pass. if self.mamba_slot_allocator is not None: self.mamba_slot_allocator.compute_and_store_offsets( req, @@ -2488,13 +2910,13 @@ def release_memory_blocks_from_request_indexes(self, request_indexes) -> None: if self.is_hybrid_model: self.mamba_metadata.free_slots(request_indexes) - # Clear intermediate offset entries for released requests + # Clear intermediate offset entries for released requests (CPU writes). if self.mamba_slot_allocator is not None: sa = self.mamba_slot_allocator - sa._intermediate_counts_gpu[request_indexes] = 0 - sa._intermediate_offsets_gpu[request_indexes] = 0 - sa._intermediate_block_ids_gpu[request_indexes] = -1 - sa._eos_cache_block_id_gpu[request_indexes] = -1 + sa._intermediate_counts_cpu[request_indexes] = 0 + sa._intermediate_offsets_cpu[request_indexes] = 0 + sa._intermediate_block_ids_cpu[request_indexes] = -1 + sa._eos_cache_block_id_cpu[request_indexes] = -1 def resume_paused_requests( self, active_request_count: int, newly_paused_request_ids: torch.Tensor @@ -2628,7 +3050,7 @@ def evict_overflow_paused_requests( -1, -1, dtype=paused_block_counts_cumsum.dtype, - device=torch.cuda.current_device(), + device='cpu', ) net_block_counts = paused_block_counts_cumsum - remaining_paused_request_counts evict_request_count = torch.nonzero(net_block_counts >= 0)[0].item() + 1 @@ -2636,9 +3058,7 @@ def evict_overflow_paused_requests( # Eviction index range. evict_start_idx = self.paused_request_count - evict_request_count evict_end_idx = self.paused_request_count - evict_request_idxs = torch.arange( - evict_start_idx, evict_end_idx, device=torch.cuda.current_device() - ) + evict_request_idxs = torch.arange(evict_start_idx, evict_end_idx, device='cpu') # Clone needed: subsequent release_memory_blocks_from_request_indexes and # _swap_book_keeping_tensors calls mutate self.request_ids in place. evict_request_ids = self.request_ids[evict_start_idx:evict_end_idx].clone() @@ -2653,24 +3073,24 @@ def evict_overflow_paused_requests( src_idxs = torch.arange( self.paused_request_count - evict_request_count, self.paused_request_count, - device=torch.cuda.current_device(), + device='cpu', ) dst_idxs = torch.arange( self.total_request_count - evict_request_count, self.total_request_count, - device=torch.cuda.current_device(), + device='cpu', ) else: # Swap all active requests with left-most evicted requests. src_idxs = torch.arange( self.paused_request_count - evict_request_count, self.paused_request_count - evict_request_count + active_request_count, - device=torch.cuda.current_device(), + device='cpu', ) dst_idxs = torch.arange( self.paused_request_count, self.paused_request_count + active_request_count, - device=torch.cuda.current_device(), + device='cpu', ) # Swap evicted and active requests. @@ -2746,6 +3166,14 @@ def update_requests( # active_request_count -> This corresponds to requests that have not reached EOD or max length # finished_request_count are requests that have reached the termination criterion + # Ensure all inputs are on CPU for bookkeeping operations. + if active_requests_mask.is_cuda: + active_requests_mask = active_requests_mask.cpu() + if new_tokens.is_cuda: + new_tokens = new_tokens.cpu() + if new_speculative_tokens is not None and new_speculative_tokens.is_cuda: + new_speculative_tokens = new_speculative_tokens.cpu() + self.num_prefill_requests = 0 # all turns to decode # All request that were in prefill become decode requests. # For the chunked prefill request we will overwrite this the next time add_request @@ -3050,14 +3478,14 @@ def update_requests( self.token_to_pos_ids[: self.active_token_count] = self.request_kv_length_offsets[ self.paused_request_count : self.total_request_count ].repeat_interleave(num_generated_tokens) + torch.arange( - num_generated_tokens, device=torch.cuda.current_device() + num_generated_tokens, device='cpu' ).repeat( active_request_count ) # # Token to request idx : [0, 0, 0, 1, 1, 1, 2, 2, 2 ...] self.token_to_request_idx[: self.active_token_count] = torch.arange( - self.paused_request_count, self.total_request_count, device=torch.cuda.current_device() + self.paused_request_count, self.total_request_count, device='cpu' ).repeat_interleave(num_generated_tokens) self.token_to_position_in_request[: self.active_token_count] = self.token_to_pos_ids[ @@ -3079,7 +3507,7 @@ def update_requests( raw_positions = ( old_offsets[:, None] + 1 # Offset by 1 because old_offsets points to the LAST token - + torch.arange(num_generated_tokens, device=torch.cuda.current_device())[None, :] + + torch.arange(num_generated_tokens, device='cpu')[None, :] ) # # A token crosses to the next block if its raw_position >= block_size @@ -3195,10 +3623,9 @@ def calculate_log_probs( # # active_token_ids[new_token_idx] = new_tokens # : [ 52 | 12 | 16 3 | 12 72 24 88 86 ] - active_token_ids = self.token_to_input_ids[: self.active_token_count].roll(-1, 0) - active_query_lengths = self.request_query_lengths[ - self.paused_request_count : self.total_request_count - ] + n_active = self.total_request_count - self.paused_request_count + active_token_ids = self.gpu_view.token_to_input_ids[: self.active_token_count].roll(-1, 0) + active_query_lengths = self.gpu_view.request_query_lengths[:n_active] new_token_idx = active_query_lengths.cumsum(0) - 1 active_token_ids[new_token_idx] = new_tokens diff --git a/megatron/core/inference/contexts/gpu_view.py b/megatron/core/inference/contexts/gpu_view.py new file mode 100644 index 00000000000..17d6e19ea03 --- /dev/null +++ b/megatron/core/inference/contexts/gpu_view.py @@ -0,0 +1,209 @@ +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +import torch + + +class ContextGPUView: + """GPU-resident snapshot of context bookkeeping data for the forward pass. + + This is the ONLY interface GPU code (attention kernels, KV append, RoPE, + sampling, log-probs, speculative verification) uses to read context state. + CPU bookkeeping code accesses context tensors directly. + + Populated once per step by ``DynamicInferenceContext.transfer_bookkeeping_to_gpu()``. + All tensors have fixed addresses for CUDA graph compatibility. + + Convention: + ``context.foo`` -> CPU (source of truth, used by bookkeeping) + ``context.gpu_view.foo`` -> GPU (snapshot, used by forward pass) + + Layout note: the 9 bookkeeping fields are backed by a single contiguous + ``uint8`` buffer (``self._buf``). Each field is a ``view(dtype)`` onto a + slice of that buffer. This matches the pinned-CPU-buffer layout in + :class:`DynamicInferenceContext` so that the per-step H2D transfer is a + single ``cudaMemcpyAsync`` instead of nine small ones. + """ + + def __init__( + self, + max_requests: int, + max_tokens: int, + max_kv_blocks: int, + device: torch.device, + max_mamba_chunks: int = 0, + ): + # Field layout (must match DynamicInferenceContext's CPU buffer layout): + # int64 token fields first (auto 8-byte alignment), then int32 token + # fields, then int32 request fields, then int32 MHA fields, then + # int32 Mamba fields (hybrid models only; omitted when + # max_mamba_chunks == 0). + tok_int64_bytes = max_tokens * 8 # 2 fields of int64 = 8 bytes/elem + tok_int32_bytes = max_tokens * 4 # 4 fields of int32 = 4 bytes/elem + req_int32_bytes = max_requests * 4 # 3 fields of int32 + + # MHA section: 5 fields shared by both graphed and non-graphed MHAMetadata + # (only one is active per step, so sharing storage is fine). + # mha_query_lengths int32 (max_bs,) = max_bs * 4 + # mha_cu_query_seq_lengths int32 (max_bs + 1,) = (max_bs+1) * 4 + # mha_kv_seq_lengths int32 (max_bs,) = max_bs * 4 + # mha_cu_kv_seq_lengths int32 (max_bs + 1,) = (max_bs+1) * 4 + # mha_block_table int32 (max_bs, max_kv_blocks) + # max_bs == max_requests in DynamicInferenceContext. + max_bs = max_requests + mha_query_lengths_bytes = max_bs * 4 + mha_cu_query_seq_lengths_bytes = (max_bs + 1) * 4 + mha_kv_seq_lengths_bytes = max_bs * 4 + mha_cu_kv_seq_lengths_bytes = (max_bs + 1) * 4 + mha_block_table_bytes = max_bs * max_kv_blocks * 4 + + # Mamba section: 9 int32 fields, only present for hybrid models. + # mamba_batch_indices_decode int32 (max_bs,) + # mamba_batch_indices_prefill int32 (max_bs,) + # mamba_seq_idx int32 (1, max_tokens) + # mamba_cu_seqlens int32 (max_bs + 1,) + # mamba_cu_chunk_seqlens int32 (max_mamba_chunks + 1,) + # mamba_last_chunk_indices int32 (max_bs,) + # mamba_seq_idx_for_varlen int32 (max_mamba_chunks,) + # mamba_conv_seq_idx int32 (max_tokens,) + # mamba_conv_seq_start int32 (max_tokens,) + if max_mamba_chunks > 0: + mamba_batch_indices_decode_bytes = max_bs * 4 + mamba_batch_indices_prefill_bytes = max_bs * 4 + mamba_seq_idx_bytes = max_tokens * 4 + mamba_cu_seqlens_bytes = (max_bs + 1) * 4 + mamba_cu_chunk_seqlens_bytes = (max_mamba_chunks + 1) * 4 + mamba_last_chunk_indices_bytes = max_bs * 4 + mamba_seq_idx_for_varlen_bytes = max_mamba_chunks * 4 + mamba_conv_seq_idx_bytes = max_tokens * 4 + mamba_conv_seq_start_bytes = max_tokens * 4 + else: + mamba_batch_indices_decode_bytes = 0 + mamba_batch_indices_prefill_bytes = 0 + mamba_seq_idx_bytes = 0 + mamba_cu_seqlens_bytes = 0 + mamba_cu_chunk_seqlens_bytes = 0 + mamba_last_chunk_indices_bytes = 0 + mamba_seq_idx_for_varlen_bytes = 0 + mamba_conv_seq_idx_bytes = 0 + mamba_conv_seq_start_bytes = 0 + + total_bytes = ( + 2 * tok_int64_bytes + + 4 * tok_int32_bytes + + 3 * req_int32_bytes + + mha_query_lengths_bytes + + mha_cu_query_seq_lengths_bytes + + mha_kv_seq_lengths_bytes + + mha_cu_kv_seq_lengths_bytes + + mha_block_table_bytes + + mamba_batch_indices_decode_bytes + + mamba_batch_indices_prefill_bytes + + mamba_seq_idx_bytes + + mamba_cu_seqlens_bytes + + mamba_cu_chunk_seqlens_bytes + + mamba_last_chunk_indices_bytes + + mamba_seq_idx_for_varlen_bytes + + mamba_conv_seq_idx_bytes + + mamba_conv_seq_start_bytes + ) + + # Zero-initialized so pre-transfer reads see zeros (matches prior semantics). + self._buf = torch.zeros(total_bytes, dtype=torch.uint8, device=device) + + # Token-level tensors (consumed by embedding, RoPE, KV append, Mamba). + off = 0 + self.token_to_input_ids = self._buf[off : off + tok_int64_bytes].view(torch.long) + off += tok_int64_bytes + self.token_to_pos_ids = self._buf[off : off + tok_int64_bytes].view(torch.long) + off += tok_int64_bytes + self.token_to_block_idx = self._buf[off : off + tok_int32_bytes].view(torch.int32) + off += tok_int32_bytes + self.token_to_local_position_within_kv_block = self._buf[off : off + tok_int32_bytes].view( + torch.int32 + ) + off += tok_int32_bytes + self.token_to_request_idx = self._buf[off : off + tok_int32_bytes].view(torch.int32) + off += tok_int32_bytes + self.token_to_position_in_request = self._buf[off : off + tok_int32_bytes].view(torch.int32) + off += tok_int32_bytes + + # Request-level tensors (consumed by sampling, log-probs, speculative verification, MTP). + self.request_in_prefill_status = self._buf[off : off + req_int32_bytes].view(torch.int32) + off += req_int32_bytes + self.request_query_lengths = self._buf[off : off + req_int32_bytes].view(torch.int32) + off += req_int32_bytes + self.request_kv_length_offsets = self._buf[off : off + req_int32_bytes].view(torch.int32) + off += req_int32_bytes + + # MHA flash-attention metadata (shared between GraphedMHAMetadata and + # NonGraphedMHAMetadata — only one is active per step). + self.mha_query_lengths = self._buf[off : off + mha_query_lengths_bytes].view(torch.int32) + off += mha_query_lengths_bytes + self.mha_cu_query_seq_lengths = self._buf[off : off + mha_cu_query_seq_lengths_bytes].view( + torch.int32 + ) + off += mha_cu_query_seq_lengths_bytes + self.mha_kv_seq_lengths = self._buf[off : off + mha_kv_seq_lengths_bytes].view(torch.int32) + off += mha_kv_seq_lengths_bytes + self.mha_cu_kv_seq_lengths = self._buf[off : off + mha_cu_kv_seq_lengths_bytes].view( + torch.int32 + ) + off += mha_cu_kv_seq_lengths_bytes + self.mha_block_table = ( + self._buf[off : off + mha_block_table_bytes] + .view(torch.int32) + .view(max_bs, max_kv_blocks) + ) + off += mha_block_table_bytes + + # Mamba varlen metadata (hybrid models only). Each GPU view matches a + # pinned CPU view in DynamicInferenceContext._cpu_bookkeeping_buf; the + # per-step coalesced H2D copy covers both MHA and Mamba alongside the + # token/request bookkeeping. + if max_mamba_chunks > 0: + self.mamba_batch_indices_decode = self._buf[ + off : off + mamba_batch_indices_decode_bytes + ].view(torch.int32) + off += mamba_batch_indices_decode_bytes + self.mamba_batch_indices_prefill = self._buf[ + off : off + mamba_batch_indices_prefill_bytes + ].view(torch.int32) + off += mamba_batch_indices_prefill_bytes + self.mamba_seq_idx = ( + self._buf[off : off + mamba_seq_idx_bytes].view(torch.int32).view(1, max_tokens) + ) + off += mamba_seq_idx_bytes + self.mamba_cu_seqlens = self._buf[off : off + mamba_cu_seqlens_bytes].view(torch.int32) + off += mamba_cu_seqlens_bytes + self.mamba_cu_chunk_seqlens = self._buf[off : off + mamba_cu_chunk_seqlens_bytes].view( + torch.int32 + ) + off += mamba_cu_chunk_seqlens_bytes + self.mamba_last_chunk_indices = self._buf[ + off : off + mamba_last_chunk_indices_bytes + ].view(torch.int32) + off += mamba_last_chunk_indices_bytes + self.mamba_seq_idx_for_varlen = self._buf[ + off : off + mamba_seq_idx_for_varlen_bytes + ].view(torch.int32) + off += mamba_seq_idx_for_varlen_bytes + self.mamba_conv_seq_idx = self._buf[off : off + mamba_conv_seq_idx_bytes].view( + torch.int32 + ) + off += mamba_conv_seq_idx_bytes + self.mamba_conv_seq_start = self._buf[off : off + mamba_conv_seq_start_bytes].view( + torch.int32 + ) + off += mamba_conv_seq_start_bytes + else: + self.mamba_batch_indices_decode = None + self.mamba_batch_indices_prefill = None + self.mamba_seq_idx = None + self.mamba_cu_seqlens = None + self.mamba_cu_chunk_seqlens = None + self.mamba_last_chunk_indices = None + self.mamba_seq_idx_for_varlen = None + self.mamba_conv_seq_idx = None + self.mamba_conv_seq_start = None + + assert off == total_bytes, f"layout bug: wrote {off} of {total_bytes} bytes" diff --git a/megatron/core/inference/contexts/kv_block_allocator.py b/megatron/core/inference/contexts/kv_block_allocator.py index 4f2744d3582..d555c925c93 100644 --- a/megatron/core/inference/contexts/kv_block_allocator.py +++ b/megatron/core/inference/contexts/kv_block_allocator.py @@ -48,30 +48,26 @@ def __init__( assert self.active_count >= 1 # ensures paused_count < total_count - 1 self.dummy_block_idx = self.total_count - 1 - # Initialize block pool as a "stack" data structure - self.block_bag = torch.arange( - self.total_count, dtype=torch.int32, device=torch.cuda.current_device() - ) + # Initialize block pool as a "stack" data structure (CPU for bookkeeping). + self.block_bag = torch.arange(self.total_count, dtype=torch.int32, device='cpu') if self.enable_prefix_caching: # Block hash tracking for prefix caching: -1 = uncomputed, positive = valid hash - self.block_hashes = torch.full( - (self.total_count,), -1, dtype=torch.int64, device=torch.cuda.current_device() - ) + self.block_hashes = torch.full((self.total_count,), -1, dtype=torch.int64, device='cpu') # Hash-to-block mapping for O(1) prefix lookup self.kv_hash_to_block_id: Dict[int, int] = {} # Reference count per block: 0 = cached (evictable), >0 = actively used self.block_ref_counts = torch.zeros( - (self.total_count,), dtype=torch.int32, device=torch.cuda.current_device() + (self.total_count,), dtype=torch.int32, device='cpu' ) # LRU timestamps for eviction ordering (higher = more recently used) # Only needed in LRU mode; RZ mode evicts immediately on ref_count==0 if self.prefix_caching_eviction_policy == PrefixCachingEvictionPolicy.LRU: self.block_timestamps = torch.zeros( - (self.total_count,), dtype=torch.int64, device=torch.cuda.current_device() + (self.total_count,), dtype=torch.int64, device='cpu' ) # Per-block MoE routing storage (populated when routing replay is enabled) @@ -247,9 +243,7 @@ def reset(self) -> None: # Without resetting the block bag, context request memory will clash and # requests will point to each other's memory blocks, resulting in faulty # generations. - self.block_bag = torch.arange( - self.total_count, dtype=torch.int32, device=torch.cuda.current_device() - ) + self.block_bag = torch.arange(self.total_count, dtype=torch.int32, device='cpu') self.total_avail = self.total_count - 1 diff --git a/megatron/core/inference/contexts/mamba_slot_allocator.py b/megatron/core/inference/contexts/mamba_slot_allocator.py index d7c57046c8a..60c8dd3416b 100644 --- a/megatron/core/inference/contexts/mamba_slot_allocator.py +++ b/megatron/core/inference/contexts/mamba_slot_allocator.py @@ -47,59 +47,70 @@ def __init__( self.max_slots = max_slots self.num_mamba_layers = num_mamba_layers - device = torch.cuda.current_device() + gpu_device = torch.cuda.current_device() num_blocks = context.kv_block_allocator.total_count - # Block <-> slot mappings - self.block_to_slot = torch.full((num_blocks,), -1, dtype=torch.int32, device=device) - self.slot_to_block = torch.full((max_slots,), -1, dtype=torch.int32, device=device) + # Block <-> slot mappings (CPU for bookkeeping). + self.block_to_slot = torch.full((num_blocks,), -1, dtype=torch.int32, device='cpu') + self.slot_to_block = torch.full((max_slots,), -1, dtype=torch.int32, device='cpu') - # Free slot pool (stack) - self.free_slots = torch.arange(max_slots, dtype=torch.int32, device=device) + # Free slot pool (stack, CPU). + self.free_slots = torch.arange(max_slots, dtype=torch.int32, device='cpu') self.free_count = max_slots - # State tensors + # State tensors (GPU - accessed by Mamba CUDA kernels). self.conv_states = torch.zeros( (num_mamba_layers, max_slots) + conv_states_shape, dtype=conv_states_dtype, - device=device, + device=gpu_device, ) self.ssm_states = torch.zeros( - (num_mamba_layers, max_slots) + ssm_states_shape, dtype=ssm_states_dtype, device=device + (num_mamba_layers, max_slots) + ssm_states_shape, + dtype=ssm_states_dtype, + device=gpu_device, ) # Hash-to-block mapping: only blocks with cached Mamba state self.hash_to_block_id: Dict[int, int] = {} - # Per-request intermediate state storage (GPU tensors, fixed-size per request) - # 0 = no offset, -1 = no block + # Per-request intermediate state storage. + # offsets_cpu and counts_cpu: CPU source of truth. GPU copies are + # populated by transfer_bookkeeping_to_gpu() since Triton kernels read them. + # block_ids and eos_cache_block_id: CPU only (consumed by CPU code). k = MAX_INTERMEDIATE_OFFSETS_PER_REQUEST - self._intermediate_offsets_gpu = torch.zeros( - (context.max_requests, k), dtype=torch.int32, device=device + self._intermediate_offsets_cpu = torch.zeros( + (context.max_requests, k), dtype=torch.int32, device='cpu' ) - self._intermediate_block_ids_gpu = torch.full( - (context.max_requests, k), -1, dtype=torch.int32, device=device + self._intermediate_counts_cpu = torch.zeros( + context.max_requests, dtype=torch.int32, device='cpu' + ) + self._intermediate_offsets_gpu = torch.zeros( + (context.max_requests, k), dtype=torch.int32, device=gpu_device ) self._intermediate_counts_gpu = torch.zeros( - context.max_requests, dtype=torch.int32, device=device + context.max_requests, dtype=torch.int32, device=gpu_device ) - self._eos_cache_block_id_gpu = torch.full( - (context.max_requests,), -1, dtype=torch.int32, device=device + # CPU-only: consumed by _collect_commit_data() which needs .tolist() anyway. + self._intermediate_block_ids_cpu = torch.full( + (context.max_requests, k), -1, dtype=torch.int32, device='cpu' + ) + self._eos_cache_block_id_cpu = torch.full( + (context.max_requests,), -1, dtype=torch.int32, device='cpu' ) # CPU flag to skip GPU sync when no intermediates exist self._has_intermediates = False - # Pre-allocated output buffers for CUDA graph compatible extraction + # Pre-allocated output buffers for CUDA graph compatible extraction (GPU). self.max_intermediate_count = MAX_INTERMEDIATE_OFFSETS_PER_REQUEST * context.max_requests self.intermediate_ssm_out = torch.zeros( (num_mamba_layers, self.max_intermediate_count) + ssm_states_shape, dtype=ssm_states_dtype, - device=device, + device=gpu_device, ) self.intermediate_conv_out = torch.zeros( (num_mamba_layers, self.max_intermediate_count) + conv_states_shape, dtype=conv_states_dtype, - device=device, + device=gpu_device, ) # ========================================================================= @@ -320,9 +331,11 @@ def store_from_live_batch(self, slots: list, request_indices: list) -> None: return device = self.conv_states.device slot_tensor = torch.tensor(slots, dtype=torch.int64, device=device) - req_tensor = torch.tensor(request_indices, dtype=torch.int64, device=device) - # Batch lookup mamba state indices (1 GPU sync) - mamba_indices = self.context.mamba_metadata.request_to_mamba_state_idx[req_tensor].tolist() + # Lookup mamba indices from CPU bookkeeping, then move to GPU for state copy. + req_tensor_cpu = torch.tensor(request_indices, dtype=torch.int64) + mamba_indices = self.context.mamba_metadata.request_to_mamba_state_idx[ + req_tensor_cpu + ].tolist() mamba_idx_tensor = torch.tensor(mamba_indices, dtype=torch.int64, device=device) # Fancy-indexed copy (2 kernel launches instead of 2E) self.conv_states[:, slot_tensor] = self.context.mamba_conv_states[:, mamba_idx_tensor] @@ -413,42 +426,39 @@ def compute_and_store_offsets( offsets = sorted(offsets_set) count = len(offsets) - # Vectorized block ID lookup: GPU gather avoids per-block .item() syncs + # CPU bookkeeping writes (no GPU kernel launches). if count > 0: - device = self._intermediate_offsets_gpu.device - abs_tokens = torch.tensor( - [skip_tokens + o for o in offsets], dtype=torch.int64, device=device - ) - block_indices = abs_tokens // ctx.block_size_tokens - 1 - bids = ctx.request_to_kv_block_ids[current_id][block_indices] + abs_tokens_cpu = torch.tensor([skip_tokens + o for o in offsets], dtype=torch.int64) + block_indices_cpu = abs_tokens_cpu // ctx.block_size_tokens - 1 + bids_cpu = ctx.request_to_kv_block_ids[current_id][block_indices_cpu] - self._intermediate_offsets_gpu[current_id, :count] = torch.tensor( - offsets, dtype=torch.int32, device=device + self._intermediate_offsets_cpu[current_id, :count] = torch.tensor( + offsets, dtype=torch.int32 ) - self._intermediate_block_ids_gpu[current_id, :count] = bids.to(torch.int32) + self._intermediate_block_ids_cpu[current_id, :count] = bids_cpu.to(torch.int32) self._has_intermediates = True - self._intermediate_counts_gpu[current_id] = count + self._intermediate_counts_cpu[current_id] = count # Block-aligned EOS: prompt_len is exactly block-aligned if last_aligned_abs == prompt_len and prompt_len > 0: last_block_idx = prompt_len // ctx.block_size_tokens - 1 if last_block_idx >= 0: - self._eos_cache_block_id_gpu[current_id] = ctx.request_to_kv_block_ids[current_id][ + self._eos_cache_block_id_cpu[current_id] = ctx.request_to_kv_block_ids[current_id][ last_block_idx ] self._has_intermediates = True else: - self._eos_cache_block_id_gpu[current_id] = -1 + self._eos_cache_block_id_cpu[current_id] = -1 else: - self._eos_cache_block_id_gpu[current_id] = -1 + self._eos_cache_block_id_cpu[current_id] = -1 - def get_intermediate_gpu_data(self): - """Get intermediate offsets and counts as GPU tensor slices for current prefill batch. + def get_intermediate_cpu_data(self): + """Get intermediate offsets and counts as CPU tensor slices for current prefill batch. Returns: - Tuple of (offsets_gpu, counts_gpu) where: - offsets_gpu: [prefill_count, 3] int32 GPU tensor - counts_gpu: [prefill_count] int32 GPU tensor + Tuple of (offsets_cpu, counts_cpu) where: + offsets_cpu: [prefill_count, 3] int32 CPU tensor + counts_cpu: [prefill_count] int32 CPU tensor Returns (None, None) if no prefill requests or no intermediates. """ if not self._has_intermediates: @@ -463,10 +473,25 @@ def get_intermediate_gpu_data(self): decode_count = ctx.batch_dimensions.decode_req_count prefill_start = active_start + decode_count - offsets = self._intermediate_offsets_gpu[prefill_start : prefill_start + prefill_count] - counts = self._intermediate_counts_gpu[prefill_start : prefill_start + prefill_count] + offsets = self._intermediate_offsets_cpu[prefill_start : prefill_start + prefill_count] + counts = self._intermediate_counts_cpu[prefill_start : prefill_start + prefill_count] return offsets, counts + def transfer_intermediate_to_gpu(self, prefill_start: int, prefill_count: int): + """Copy intermediate offsets/counts slice from CPU to GPU for Mamba kernels. + + Returns the GPU tensor views for the forward-pass kernels to consume. + """ + if prefill_count == 0: + return None, None + offsets_cpu = self._intermediate_offsets_cpu[prefill_start : prefill_start + prefill_count] + counts_cpu = self._intermediate_counts_cpu[prefill_start : prefill_start + prefill_count] + offsets_gpu = self._intermediate_offsets_gpu[prefill_start : prefill_start + prefill_count] + counts_gpu = self._intermediate_counts_gpu[prefill_start : prefill_start + prefill_count] + offsets_gpu.copy_(offsets_cpu, non_blocking=True) + counts_gpu.copy_(counts_cpu, non_blocking=True) + return offsets_gpu, counts_gpu + # ========================================================================= # Intermediate state commit # ========================================================================= @@ -517,14 +542,14 @@ def _collect_commit_data(self): decode_count = ctx.batch_dimensions.decode_req_count prefill_start = active_start + decode_count - # Batch-transfer block IDs and EOS block IDs from GPU (2 GPU syncs) + # Block IDs and EOS block IDs live on CPU (no GPU sync needed). intermediate_count = metadata.intermediate_count per_request_counts = metadata.per_request_intermediate_counts - all_block_ids_cpu = self._intermediate_block_ids_gpu[ + all_block_ids_cpu = self._intermediate_block_ids_cpu[ prefill_start : prefill_start + prefill_count ].tolist() - eos_bids_cpu = self._eos_cache_block_id_gpu[ + eos_bids_cpu = self._eos_cache_block_id_cpu[ prefill_start : prefill_start + prefill_count ].tolist() @@ -586,10 +611,10 @@ def _clear_intermediate_state(self) -> None: decode_count = ctx.batch_dimensions.decode_req_count prefill_start = active_start + decode_count end = prefill_start + prefill_count - self._intermediate_counts_gpu[prefill_start:end].fill_(0) - self._intermediate_offsets_gpu[prefill_start:end].fill_(0) - self._intermediate_block_ids_gpu[prefill_start:end].fill_(-1) - self._eos_cache_block_id_gpu[prefill_start:end].fill_(-1) + self._intermediate_counts_cpu[prefill_start:end].fill_(0) + self._intermediate_offsets_cpu[prefill_start:end].fill_(0) + self._intermediate_block_ids_cpu[prefill_start:end].fill_(-1) + self._eos_cache_block_id_cpu[prefill_start:end].fill_(-1) self._has_intermediates = False # ========================================================================= @@ -600,15 +625,13 @@ def reset(self) -> None: """Reset all state (mappings, free pool, cache, intermediate tracking).""" self.block_to_slot.fill_(-1) self.slot_to_block.fill_(-1) - self.free_slots = torch.arange( - self.max_slots, dtype=torch.int32, device=torch.cuda.current_device() - ) + self.free_slots = torch.arange(self.max_slots, dtype=torch.int32, device='cpu') self.free_count = self.max_slots self.hash_to_block_id.clear() self.intermediate_ssm_out.zero_() self.intermediate_conv_out.zero_() - self._intermediate_offsets_gpu.fill_(0) - self._intermediate_block_ids_gpu.fill_(-1) - self._intermediate_counts_gpu.fill_(0) - self._eos_cache_block_id_gpu.fill_(-1) + self._intermediate_offsets_cpu.fill_(0) + self._intermediate_counts_cpu.fill_(0) + self._intermediate_block_ids_cpu.fill_(-1) + self._eos_cache_block_id_cpu.fill_(-1) self._has_intermediates = False diff --git a/megatron/core/inference/engines/async_zmq_communicator.py b/megatron/core/inference/engines/async_zmq_communicator.py index 52570845d61..aa13f659d40 100644 --- a/megatron/core/inference/engines/async_zmq_communicator.py +++ b/megatron/core/inference/engines/async_zmq_communicator.py @@ -131,6 +131,45 @@ async def all_reduce_max(self, *local_vals: int, async_op=True) -> int | tuple[i except zmq.Again: await asyncio.sleep(0.001) + def sync_all_reduce_max(self, *local_vals: int) -> int | tuple[int, ...]: + """Synchronous (non-asyncio) variant of all_reduce_max. + + Uses blocking ZMQ sends/recvs so it can be called from synchronous + call sites that need a CPU-only MAX reduction across the process + group. Intended for tiny payloads (e.g. a few integers) that would + otherwise force a NCCL AllReduce kernel on the compute stream. + + Note: when called from inside a running asyncio event loop, the + blocking recv will pause other coroutines on this rank until all + peers respond. This is acceptable here because every rank reaches + the call simultaneously and the message size is trivial. + + Returns a single int when called with one argument, otherwise a tuple. + """ + n = len(local_vals) + if n == 0: + raise ValueError("sync_all_reduce_max requires at least one value") + + if self.world_size <= 1: + return local_vals[0] if n == 1 else local_vals + + fmt = f'!{n}i' + payload = struct.pack(fmt, *local_vals) + + if self.is_leader: + rows = [local_vals] + while len(rows) < self.world_size: + msg = self.gather_sock.recv() + rows.append(struct.unpack(fmt, msg)) + maxes = tuple(max(row[i] for row in rows) for i in range(n)) + self.bcast_sock.send(struct.pack(fmt, *maxes)) + return maxes[0] if n == 1 else maxes + else: + self.gather_sock.send(payload) + msg = self.bcast_sock.recv() + result = struct.unpack(fmt, msg) + return result[0] if n == 1 else result + def close(self): """ Close the ZMQ sockets. diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index 64356e98034..ff5454bbaa3 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -639,6 +639,10 @@ async def start_listening_to_data_parallel_coordinator( self.expert_parallel_zmq_communicator = AsyncZMQCommunicator( self.zmq_context, process_group=self.pg_collection.ep, hostname=hostname ) + # Give the context a CPU-side MAX-reduction primitive so + # match_graph_config() can avoid a per-step NCCL AllReduce kernel. + if hasattr(self.context, "set_ep_zmq_communicator"): + self.context.set_ep_zmq_communicator(self.expert_parallel_zmq_communicator) # initialize zmq-based world communicator for consensus barriers total_world_size = torch.distributed.get_world_size() @@ -1189,10 +1193,15 @@ def post_process_requests( request.ttft = ( first_token_event.timestamp - request.event_add_engine.timestamp ) - if request.tpot is None: - request.tpot = [] - per_token_step_time = step_time / len(tokens) - request.tpot.extend([per_token_step_time] * len(tokens)) + # TPOT is observability-only. step_time is 0.0 on + # non-logging steps (async_forward skips the event sync), + # so gate the update to keep the metric a truthful sparse + # sample instead of polluting it with zeros. + if step_time > 0: + if request.tpot is None: + request.tpot = [] + per_token_step_time = step_time / len(tokens) + request.tpot.extend([per_token_step_time] * len(tokens)) # Check for stop words (after token is appended). # With speculative decoding, a stop word may end before the last @@ -1671,56 +1680,74 @@ async def async_forward(self) -> Tuple[Dict, Dict, float]: # schedule requests self.schedule_waiting_requests() - # Saving pre-step state, for printing output below. + # The print block (async_bookkeep) and metrics block both fire on this + # condition after step_count is incremented. Predict it up-front so we + # can skip the GPU-timing sync and the context_state dict builds that + # only exist to feed those logging/metrics blocks. + will_log_this_step = ( + self.logging_step_interval > 0 + and (self.context.step_count + 1) % self.logging_step_interval == 0 + ) + is_decode_only = self.context.is_decode_only() - pre_step_context_state = { - "is_decode_only": is_decode_only, - "max_requests": self.context.max_requests, - "total_request_count": self.context.total_request_count, - "paused_request_count": self.context.paused_request_count, - "active_token_count": self.context.active_token_count, - "step_count": self.context.step_count, - } + if will_log_this_step: + pre_step_context_state = { + "is_decode_only": is_decode_only, + "max_requests": self.context.max_requests, + "total_request_count": self.context.total_request_count, + "paused_request_count": self.context.paused_request_count, + "active_token_count": self.context.active_token_count, + "step_count": self.context.step_count, + } + else: + # active_token_count and step_count are still consumed by + # post_process_requests' pre_fwd_* args (for add_event_generated_token); + # the other four fields are only read in the gated print block. + pre_step_context_state = { + "active_token_count": self.context.active_token_count, + "step_count": self.context.step_count, + } # Generate tokens. nvtx_range_push("Prefill" if not is_decode_only else "Decode") # TODO @TDE: Account for this line when overlapping forward and bookkeep. self.is_decode_only = is_decode_only - self.step_start_event.record() + if will_log_this_step: + self.step_start_event.record() result = await self.controller.async_generate_output_tokens_dynamic_batch() - self.step_end_event.record() - self.step_end_event.synchronize() - step_time = self.step_start_event.elapsed_time(self.step_end_event) / 1e3 + if will_log_this_step: + self.step_end_event.record() + self.step_end_event.synchronize() + step_time = self.step_start_event.elapsed_time(self.step_end_event) / 1e3 + else: + step_time = 0.0 self.context.step_count += 1 self.context.prefix_cache_lru_clock += 1 nvtx_range_pop("Prefill" if not is_decode_only else "Decode") - if ( - self.logging_step_interval > 0 - and self.context.step_count > 0 - and self.context.step_count % self.logging_step_interval == 0 - and self.metrics_writer is not None - ): - kvcache_util_stats = self.context.get_kvcache_utilization_stats() + if will_log_this_step: + kvcache_util_stats = ( + self.context.get_kvcache_utilization_stats() + if self.metrics_writer is not None + else None + ) + post_step_context_state = { + "waiting_request_count": len(self.waiting_request_ids), + "finished_request_count": self.finished_request_count, + "evicted_request_count": self.evicted_request_count, + "kv_stats": kvcache_util_stats, + "total_active_block_count": self.context.kv_block_allocator.active_count, + "total_paused_block_count": self.context.kv_block_allocator.paused_count, + "total_active_used_blocks": self.context.kv_block_allocator.get_active_used(), + "total_paused_used_blocks": self.context.kv_block_allocator.get_paused_used(), + } + context_state = {**pre_step_context_state, **post_step_context_state} else: - kvcache_util_stats = None - - post_step_context_state = { - "waiting_request_count": len(self.waiting_request_ids), - "finished_request_count": self.finished_request_count, - "evicted_request_count": self.evicted_request_count, - "kv_stats": kvcache_util_stats, - "padded_active_token_count": self.context.padded_active_token_count, - "using_cuda_graph_this_step": self.context.using_cuda_graph_this_step(), - "total_active_block_count": self.context.kv_block_allocator.active_count, - "total_paused_block_count": self.context.kv_block_allocator.paused_count, - "total_active_used_blocks": self.context.kv_block_allocator.get_active_used(), - "total_paused_used_blocks": self.context.kv_block_allocator.get_paused_used(), - } - - context_state = {**pre_step_context_state, **post_step_context_state} + # Keep kv_stats=None so the metrics-block gate at `async_bookkeep` + # (`if context_state["kv_stats"] is not None`) remains well-typed. + context_state = {**pre_step_context_state, "kv_stats": None} return result, context_state, step_time diff --git a/megatron/core/inference/text_generation_controllers/mtp_utils_pytorch.py b/megatron/core/inference/text_generation_controllers/mtp_utils_pytorch.py index 59bad67d70a..fe5474d0b22 100644 --- a/megatron/core/inference/text_generation_controllers/mtp_utils_pytorch.py +++ b/megatron/core/inference/text_generation_controllers/mtp_utils_pytorch.py @@ -36,6 +36,17 @@ def rewind_kv_cache( if num_active_requests is None: num_active_requests = N + # Bulk-extract scalars once via .tolist() instead of per-element .item(). + # Avoids N round-trips through the Python/C++ boundary inside the loop. + accepted_list = accepted_counts.tolist() + prefill_list = prefill_status.tolist() + offset_list = last_kv_block_offset.tolist() + length_list = kv_length_offsets.tolist() + block_count_list = kv_block_counts.tolist() + last_block_list = last_kv_block_id.tolist() + kv_block_ids_list = kv_block_ids.tolist() + max_blocks = kv_block_ids.shape[1] + blocks_to_release = torch.empty_like(last_kv_block_id) remove_mask = torch.empty(N, device=accepted_counts.device, dtype=torch.bool) @@ -45,12 +56,12 @@ def rewind_kv_cache( remove_mask[i] = False continue - accepted = accepted_counts[i].item() - prefill = prefill_status[i].item() - last_offset = last_kv_block_offset[i].item() - kv_length = kv_length_offsets[i].item() - block_count = kv_block_counts[i].item() - last_block = last_kv_block_id[i].item() + accepted = accepted_list[i] + prefill = prefill_list[i] + last_offset = offset_list[i] + kv_length = length_list[i] + block_count = block_count_list[i] + last_block = last_block_list[i] # Number of tokens to rewind (rejected speculative tokens). # For prefill requests, no speculative tokens were forwarded through the model, @@ -81,11 +92,11 @@ def rewind_kv_cache( # Update last_kv_block_id to point to the previous block (at index new_count - 1) prev_idx = max(new_block_count - 1, 0) - prev_block_id = kv_block_ids[i, prev_idx].item() + prev_block_id = kv_block_ids_list[i][prev_idx] last_kv_block_id[i] = prev_block_id if remove else last_block # Clear the released block entry (at index new_count, which was the old last block) - scatter_idx = min(new_block_count, kv_block_ids.shape[1] - 1) + scatter_idx = min(new_block_count, max_blocks - 1) if remove: kv_block_ids[i, scatter_idx] = -1 diff --git a/megatron/core/inference/text_generation_controllers/text_generation_controller.py b/megatron/core/inference/text_generation_controllers/text_generation_controller.py index d665e17dc1a..faa9e5babd6 100644 --- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py +++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py @@ -11,6 +11,7 @@ import torch import torch.nn.functional as F from torch import Tensor +from torch.cuda.nvtx import range_pop, range_push from megatron.core import parallel_state from megatron.core.inference.async_stream import AsyncStream @@ -59,10 +60,10 @@ HAVE_TE = False from megatron.core.inference.batch_dimensions_utils import InferenceBatchDimensions +from megatron.core.inference.text_generation_controllers.mtp_utils_pytorch import rewind_kv_cache from megatron.core.inference.text_generation_controllers.mtp_utils_triton import ( mamba_state_selective_copy, prepare_next_forward_pass, - rewind_kv_cache, verify_speculative_tokens, ) @@ -612,11 +613,19 @@ def _dynamic_step_context_init( unwrapped_model = unwrap_model(self.inference_wrapped_model.model) model_config = get_model_config(unwrapped_model) - # Initialize attention state. + # Initialize attention state (100% CPU computation). + range_push("initialize_attention_state") context.initialize_attention_state( construct_graph_dimensions=construct_graph_dimensions, is_expert_parallel_dummy_cuda_graph_step=is_dummy_forward, ) + range_pop() + + # Single batch CPU-to-GPU transfer of bookkeeping state. + range_push("transfer_bookkeeping_to_gpu") + context.transfer_bookkeeping_to_gpu() + range_pop() + set_moe_metadata_sync(unwrapped_model) # Derive the MTP padded batch size from the existing padded graph dimensions. @@ -756,9 +765,10 @@ def _rewind_kv_cache(self) -> tuple: """Update the KV cache bookkeeping for speculative decoding. After forward pass with speculative tokens, some tokens may be rejected. - This function "rewinds" the KV cache bookkeeping to reflect only the accepted - tokens. The core bookkeeping is handled by a Triton kernel (one thread per - request). Mamba hybrid-model state updates remain in PyTorch. + This function "rewinds" the KV cache bookkeeping to reflect only the + accepted tokens. The core bookkeeping rewind runs on CPU (mutating the + CPU source-of-truth tensors in place); the Mamba hybrid-model state + update stays on GPU because it operates on GPU-resident state buffers. Returns (blocks_to_release, remove_mask) for the caller to release blocks back to the allocator outside the compiled graph. @@ -767,48 +777,51 @@ def _rewind_kv_cache(self) -> tuple: active_request_count = context.total_request_count - context.paused_request_count active_request_slice = slice(context.paused_request_count, context.total_request_count) - accepted_tokens_per_request = self._accepted_token_counts_per_request[:active_request_count] - - request_in_prefill_status = context.request_in_prefill_status_tensor[active_request_slice] - request_last_kv_block_offset = context.request_last_kv_block_offset[active_request_slice] - request_kv_length_offsets = context.request_kv_length_offsets[active_request_slice] - request_kv_block_counts = context.request_kv_block_counts[active_request_slice] - request_last_kv_block_id = context.request_last_kv_block_id[active_request_slice] - request_to_kv_block_ids = context.request_to_kv_block_ids[active_request_slice] + # accepted_counts is the only GPU input; D2H a small slice so the + # CPU rewind can read its values via .tolist() inside a Python loop. + accepted_tokens_per_request_cpu = self._accepted_token_counts_per_request[ + :active_request_count + ].cpu() - # --- Triton kernel: core KV-cache rewind --- blocks_to_release, remove_mask = rewind_kv_cache( - accepted_counts=accepted_tokens_per_request, - prefill_status=request_in_prefill_status, - last_kv_block_offset=request_last_kv_block_offset, - kv_length_offsets=request_kv_length_offsets, - kv_block_counts=request_kv_block_counts, - last_kv_block_id=request_last_kv_block_id, - kv_block_ids=request_to_kv_block_ids, + accepted_counts=accepted_tokens_per_request_cpu, + prefill_status=context.request_in_prefill_status_tensor[active_request_slice], + last_kv_block_offset=context.request_last_kv_block_offset[active_request_slice], + kv_length_offsets=context.request_kv_length_offsets[active_request_slice], + kv_block_counts=context.request_kv_block_counts[active_request_slice], + last_kv_block_id=context.request_last_kv_block_id[active_request_slice], + kv_block_ids=context.request_to_kv_block_ids[active_request_slice], num_speculative_tokens=self.num_speculative_tokens, block_size_tokens=context.block_size_tokens, num_active_requests=active_request_count, ) - # Mamba speculative rewind: copy accepted intermediate states in-place. + # Mamba speculative rewind stays on GPU because it mutates GPU-resident + # SSM/conv state that the next forward pass reads directly. if context.is_hybrid_model: + cuda_device = torch.cuda.current_device() + # gpu_view.request_in_prefill_status was uploaded by this step's + # coalesced H2D and mirrors the active-slice CPU values, so we + # don't need to re-upload prefill_status for the Mamba kernels. + prefill_status_gpu = context.gpu_view.request_in_prefill_status[:active_request_count] + accepted_counts_gpu = self._accepted_token_counts_per_request[:active_request_count] mamba_state_idx = context.mamba_metadata.request_to_mamba_state_idx[ active_request_slice - ] + ].to(cuda_device, non_blocking=True) mamba_state_selective_copy( intermediate_states=context.mamba_intermediate_conv_states, current_states=context.mamba_conv_states, - prefill_status=request_in_prefill_status, + prefill_status=prefill_status_gpu, state_idx=mamba_state_idx, - accepted_counts=accepted_tokens_per_request, + accepted_counts=accepted_counts_gpu, num_layers=context.num_mamba_layers, ) mamba_state_selective_copy( intermediate_states=context.mamba_intermediate_ssm_states, current_states=context.mamba_ssm_states, - prefill_status=request_in_prefill_status, + prefill_status=prefill_status_gpu, state_idx=mamba_state_idx, - accepted_counts=accepted_tokens_per_request, + accepted_counts=accepted_counts_gpu, num_layers=context.num_mamba_layers, ) @@ -877,11 +890,16 @@ def _compute_serial_mtp_and_sample(self): last_accepted_hidden = None # Compute position IDs for the next tokens. - # After rewind, request_kv_length_offsets has been adjusted. The actual - # KV cache length is: adjusted_offset + processed_tokens. - # The next position to predict starts at that cache length. - adjusted_offsets = context.request_kv_length_offsets[active_slice] - processed_tokens = context.request_query_lengths[active_slice] + # After rewind, request_kv_length_offsets has been adjusted. Read from + # CPU context (post-rewind values), NOT gpu_view (stale pre-rewind snapshot). + # The next position to predict is: adjusted_offset + processed_tokens. + cuda_device = torch.cuda.current_device() + adjusted_offsets = context.request_kv_length_offsets[active_slice].to( + cuda_device, non_blocking=True + ) + processed_tokens = context.request_query_lengths[active_slice].to( + cuda_device, non_blocking=True + ) # Cast to int64 to match CUDA graph capture dtype expectations. base_position = (adjusted_offsets + processed_tokens).to(torch.int64) @@ -994,6 +1012,7 @@ def _sample_speculative_logits( Returns: tuple: (output_tokens, repeats) where output_tokens has shape [total_required_tokens] """ + # request_in_prefill_status_tensor is already on GPU (from gpu_view). repeats = torch.where( request_in_prefill_status_tensor == 0, 1 + self.num_speculative_tokens, 1 ) @@ -1053,8 +1072,9 @@ def _dynamic_step_sample_logits_and_verify_tokens(self, input_ids: Tensor): context = self.inference_wrapped_model.inference_context active_request_count = context.total_request_count - context.paused_request_count - request_in_prefill_status_tensor = context.request_in_prefill_status_tensor[ - context.paused_request_count : context.total_request_count + # Use gpu_view for data consumed by GPU operations (sampling, verification). + request_in_prefill_status_tensor = context.gpu_view.request_in_prefill_status[ + :active_request_count ] # Get the logit indices for tokens that need sampling. @@ -1294,12 +1314,11 @@ def _dynamic_step_calculate_log_probs_speculative(self) -> Tuple[List[List[float context = self.inference_wrapped_model.inference_context active_request_count = context.total_request_count - context.paused_request_count - request_in_prefill_status_tensor = context.request_in_prefill_status_tensor[ - context.paused_request_count : context.total_request_count - ] - request_query_lengths = context.request_query_lengths[ - context.paused_request_count : context.total_request_count + # Use gpu_view for data consumed by GPU log-probs operations. + request_in_prefill_status_tensor = context.gpu_view.request_in_prefill_status[ + :active_request_count ] + request_query_lengths = context.gpu_view.request_query_lengths[:active_request_count] num_prefill_requests = request_in_prefill_status_tensor.sum().item() num_decode_requests = active_request_count - num_prefill_requests @@ -1362,7 +1381,7 @@ def _dynamic_step_calculate_log_probs_speculative(self) -> Tuple[List[List[float ] log_probs_list_prefill = [[lp.item()] for lp in selected_log_probs] else: - prefill_token_ids = context.token_to_input_ids[ + prefill_token_ids = context.gpu_view.token_to_input_ids[ decode_len : context.active_token_count ].roll(-1, 0) prefill_query_lengths = request_query_lengths[request_in_prefill_status_tensor == 1] @@ -1405,12 +1424,11 @@ def _dynamic_step_calculate_top_n_logprobs_speculative( context = self.inference_wrapped_model.inference_context active_request_count = context.total_request_count - context.paused_request_count - request_in_prefill_status_tensor = context.request_in_prefill_status_tensor[ - context.paused_request_count : context.total_request_count - ] - request_query_lengths = context.request_query_lengths[ - context.paused_request_count : context.total_request_count + # Use gpu_view for data consumed by GPU top-n operations. + request_in_prefill_status_tensor = context.gpu_view.request_in_prefill_status[ + :active_request_count ] + request_query_lengths = context.gpu_view.request_query_lengths[:active_request_count] num_prefill_requests = request_in_prefill_status_tensor.sum().item() num_decode_requests = active_request_count - num_prefill_requests @@ -1701,6 +1719,24 @@ def _dummy_serial_mtp_forward(self): ) nvtx_range_pop(f"mtp-spec-decoding/dummy-depth-{depth}") + def _transfer_samples_to_cpu(self, active_request_count: int) -> tuple: + """Batch GPU-to-CPU transfer of sampled tokens. + + Called at the boundary between GPU sampling and CPU bookkeeping. + After this returns, all sampled data is on CPU and the remainder + of the step is 100% CPU. + + Returns: + tuple: (sampled_tokens_cpu, sampled_mtp_tokens_cpu) where + sampled_mtp_tokens_cpu is None when speculative decoding is off. + """ + sampled_tokens_cpu = self._sampled_tokens_cuda[:active_request_count].cpu() + if self.num_speculative_tokens > 0: + sampled_mtp_tokens_cpu = self._sampled_mtp_tokens_cuda[:, :active_request_count].cpu() + else: + sampled_mtp_tokens_cpu = None + return sampled_tokens_cpu, sampled_mtp_tokens_cpu + def _dynamic_step_context_bookkeeping(self) -> Dict[str, Tensor]: """Update the dynamic inference context after sampling. @@ -1720,7 +1756,15 @@ def _dynamic_step_context_bookkeeping(self) -> Dict[str, Tensor]: active_request_count = context.total_request_count - context.paused_request_count active_request_slice = slice(context.paused_request_count, context.total_request_count) - # Active sequence lengths. + # Batch GPU-to-CPU transfer of all sampled tokens. + range_push("transfer_samples_to_cpu") + sampled_tokens_cpu, sampled_mtp_tokens_cpu = self._transfer_samples_to_cpu( + active_request_count + ) + range_pop() + + range_push("active_request_mask") + # Everything below is 100% CPU. active_request_ids = context.request_ids[active_request_slice].long() active_sequence_lengths = context.get_active_sequence_lengths() @@ -1732,9 +1776,10 @@ def _dynamic_step_context_bookkeeping(self) -> Dict[str, Tensor]: max_sequence_lengths = context.get_max_sequence_lengths() # Request finished if termination_id or length >= max_sequence_length. - # Note: termination_id tensor has per-request termination IDs from mixed sampling + # Both operands are CPU: sampled_tokens_cpu was D2H'd above, and + # active_request_metadata is CPU-pinned. active_request_mask = ( - self._sampled_tokens_cuda[:active_request_count] + sampled_tokens_cpu != context.active_request_metadata["termination_id"][:active_request_count] ).byte() & torch.less(active_sequence_lengths, max_sequence_lengths).byte() @@ -1765,22 +1810,24 @@ def _dynamic_step_context_bookkeeping(self) -> Dict[str, Tensor]: finished_routing_block_ids[req_id] = valid # Clone needed: update_requests mutates next_tokens in-place via tensor_swap, - # which would corrupt the reused _sampled_tokens_cuda buffer. - new_sample_copy = self._sampled_tokens_cuda[:active_request_count].clone() + # which would corrupt the reused buffer. + new_sample_copy = sampled_tokens_cpu.clone() + range_pop() - # Update requests. - # _sampled_mtp_tokens_cuda has shape [num_speculative_tokens, max_requests] - if self.num_speculative_tokens > 0: - sampled_mtp_tokens_cuda = self._sampled_mtp_tokens_cuda[:, :active_request_count] - else: - sampled_mtp_tokens_cuda = None + range_push("update_requests") update_result = context.update_requests( - active_request_mask, new_sample_copy, sampled_mtp_tokens_cuda + active_request_mask, new_sample_copy, sampled_mtp_tokens_cpu ) + range_pop() return { "active_request_ids": active_request_ids, "finished_request_ids": finished_request_ids, + # Already a CPU tensor (independent of _sampled_tokens_cuda via the + # .cpu() in _transfer_samples_to_cpu; update_requests only mutates + # the separate new_sample_copy). Returning the CPU copy avoids a + # D2H sync when the engine later calls sample.tolist(). + "sample": sampled_tokens_cpu, "finished_routing_block_ids": finished_routing_block_ids, **(update_result or {}), } @@ -1825,6 +1872,7 @@ async def async_generate_output_tokens_dynamic_batch( # Forward pass produces only base logits. When speculative decoding is # active, MTP logits are computed serially after verification. + range_push("forward_pass") self._dynamic_step_forward_logits(input_ids, position_ids) # Commit Mamba intermediate states before update_requests, which @@ -1838,6 +1886,7 @@ async def async_generate_output_tokens_dynamic_batch( # Must be done before update_requests while token-to-block mappings are valid. # Reconstruction happens from blocks at request completion. context.kv_block_allocator.store_routing_per_block(self._router_record_bookkeeping()) + range_pop() # This is the best place to yield control back to event loop. # At this point we have enqueued FW pass GPU kernels asynchronously. @@ -1849,6 +1898,7 @@ async def async_generate_output_tokens_dynamic_batch( await asyncio.sleep(0) with torch.inference_mode(): + range_push("sampling") return_log_probs, return_top_n_logprobs = self._dynamic_step_log_probs_bookkeeping() self._dynamic_step_sample_bookkeeping() @@ -1897,15 +1947,21 @@ async def async_generate_output_tokens_dynamic_batch( top_n_logprobs = self._dynamic_step_calculate_top_n_logprobs( log_probs_tensor ) + range_pop() if skip_bookkeeping: - request_bookkeeping = {} + # _transfer_samples_to_cpu wasn't invoked on this path, so do + # a one-shot D2H here to keep "sample" as a CPU tensor for + # downstream consumers. + request_bookkeeping = { + "sample": self._sampled_tokens_cuda[:active_request_count].cpu() + } else: + # request_bookkeeping supplies "sample" as the already-CPU + # tensor produced by _transfer_samples_to_cpu. request_bookkeeping = self._dynamic_step_context_bookkeeping() ret = { - # Clone needed: _sampled_tokens_cuda is a reused buffer overwritten each step. - "sample": self._sampled_tokens_cuda[:active_request_count].clone(), "accepted_tokens": ( # Clone needed: .fill_(-1) on line 1480 would corrupt the returned value. self._accepted_tokens_per_request.clone() diff --git a/tests/unit_tests/inference/contexts/test_dynamic_context.py b/tests/unit_tests/inference/contexts/test_dynamic_context.py index b20686685cd..0e307f600f1 100644 --- a/tests/unit_tests/inference/contexts/test_dynamic_context.py +++ b/tests/unit_tests/inference/contexts/test_dynamic_context.py @@ -221,7 +221,7 @@ def test_request_overflow(self, is_hybrid_model: bool): dynamic_context.add_request( DynamicInferenceRequest( request_id=i, - prompt_tokens=torch.zeros(10, device='cuda'), + prompt_tokens=torch.zeros(10, device='cpu'), sampling_params=SamplingParams( num_tokens_to_generate=dynamic_context.max_tokens - 10 ), @@ -249,7 +249,7 @@ def test_token_overflow_error(self, is_hybrid_model: bool): dynamic_context.add_request( DynamicInferenceRequest( request_id=1, - prompt_tokens=torch.arange(0, 225, device='cuda'), + prompt_tokens=torch.arange(0, 225, device='cpu'), sampling_params=SamplingParams( num_tokens_to_generate=dynamic_context.max_tokens - 25 ), @@ -279,7 +279,7 @@ def test_reset(self, is_hybrid_model: bool): dynamic_context.paused_request_count = 5 dynamic_context.padded_active_token_count = 10 dynamic_context.padded_active_request_count = 5 - dynamic_context.paused_tokens = torch.tensor([1, 2, 3], device='cuda') + dynamic_context.paused_tokens = torch.tensor([1, 2, 3], device='cpu') dynamic_context.request_ids.fill_(1) dynamic_context.request_query_lengths.fill_(1) dynamic_context.request_kv_length_offsets.fill_(1) @@ -363,7 +363,7 @@ def test_allocate_and_release_memory_blocks(self, is_hybrid_model): ) assert dynamic_context.kv_block_allocator.total_avail == expected_block_count_avail dynamic_context.kv_block_allocator.release_memory_blocks( - torch.tensor(expected_memory_blocks[-2:], device='cuda') + torch.tensor(expected_memory_blocks[-2:], device='cpu') ) assert dynamic_context.kv_block_allocator.total_avail == expected_block_count_avail + 2 assert ( @@ -400,7 +400,7 @@ def test_add_request(self, is_hybrid_model: bool): dynamic_context.add_request( DynamicInferenceRequest( request_id=0, - prompt_tokens=torch.arange(0, context_length, dtype=torch.long, device='cuda'), + prompt_tokens=torch.arange(0, context_length, dtype=torch.long, device='cpu'), sampling_params=SamplingParams( num_tokens_to_generate=dynamic_context.max_tokens - context_length ), @@ -419,15 +419,15 @@ def test_add_request(self, is_hybrid_model: bool): assert dynamic_context.request_last_kv_block_offset[0].item() == 15 assert torch.all( dynamic_context.token_to_pos_ids[0:context_length] - == torch.arange(0, context_length, dtype=torch.long, device='cuda') + == torch.arange(0, context_length, dtype=torch.long, device='cpu') ) assert torch.all( dynamic_context.token_to_input_ids[0:context_length] - == torch.arange(0, context_length, dtype=torch.long, device='cuda') + == torch.arange(0, context_length, dtype=torch.long, device='cpu') ) assert torch.all( dynamic_context.token_to_position_in_request[0:context_length] - == torch.arange(0, context_length, dtype=torch.long, device='cuda') + == torch.arange(0, context_length, dtype=torch.long, device='cpu') ) # Verify token_to_block_idx and token_to_local_position_within_kv_block based on assigned blocks @@ -448,7 +448,7 @@ def test_add_request(self, is_hybrid_model: bool): ) assert torch.all( dynamic_context.token_to_local_position_within_kv_block[0:context_length] - == torch.arange(0, context_length, dtype=torch.long, device='cuda') + == torch.arange(0, context_length, dtype=torch.long, device='cpu') % dynamic_context.block_size_tokens ) @@ -470,12 +470,12 @@ def test_add_dummy_requests_parallel_populates_state(self): requests = [ DynamicInferenceRequest( request_id=100, - prompt_tokens=torch.arange(0, 3, device='cuda'), + prompt_tokens=torch.arange(0, 3, device='cpu'), sampling_params=SamplingParams(num_tokens_to_generate=2, termination_id=7), ), DynamicInferenceRequest( request_id=101, - prompt_tokens=torch.arange(3, 9, device='cuda'), + prompt_tokens=torch.arange(3, 9, device='cpu'), sampling_params=SamplingParams(num_tokens_to_generate=1, termination_id=8), ), ] @@ -492,12 +492,12 @@ def test_add_dummy_requests_parallel_populates_state(self): assert dynamic_context.kv_block_allocator.total_avail == block_avail_before expected_tokens = torch.cat( - [torch.arange(0, 3, device='cuda'), torch.arange(3, 9, device='cuda')] + [torch.arange(0, 3, device='cpu'), torch.arange(3, 9, device='cpu')] ) assert torch.equal(dynamic_context.token_to_input_ids[:total_tokens], expected_tokens) expected_positions = torch.tensor( - [0, 1, 2, 0, 1, 2, 3, 4, 5], device='cuda', dtype=torch.long + [0, 1, 2, 0, 1, 2, 3, 4, 5], device='cpu', dtype=torch.long ) assert torch.equal( dynamic_context.token_to_position_in_request[:total_tokens], expected_positions @@ -505,7 +505,7 @@ def test_add_dummy_requests_parallel_populates_state(self): assert torch.equal(dynamic_context.token_to_pos_ids[:total_tokens], expected_positions) expected_request_indices = torch.tensor( - [0, 0, 0, 1, 1, 1, 1, 1, 1], device='cuda', dtype=torch.long + [0, 0, 0, 1, 1, 1, 1, 1, 1], device='cpu', dtype=torch.long ) assert torch.equal( dynamic_context.token_to_request_idx[:total_tokens], expected_request_indices @@ -521,15 +521,15 @@ def test_add_dummy_requests_parallel_populates_state(self): assert torch.equal( dynamic_context.request_query_lengths[: len(requests)], - torch.tensor(lengths, device='cuda', dtype=torch.int32), + torch.tensor(lengths, device='cpu', dtype=torch.int32), ) assert torch.equal( dynamic_context.request_output_lengths[: len(requests)], - torch.tensor([5, 7], device='cuda', dtype=torch.int32), + torch.tensor([5, 7], device='cpu', dtype=torch.int32), ) assert torch.equal( dynamic_context.request_kv_block_counts[: len(requests)], - torch.tensor([1, 2], device='cuda', dtype=torch.int32), + torch.tensor([1, 2], device='cpu', dtype=torch.int32), ) assert torch.all( dynamic_context.request_to_kv_block_ids[0, :1] == dummy_block_idx @@ -542,12 +542,12 @@ def test_add_dummy_requests_parallel_populates_state(self): assert torch.all(dynamic_context.request_last_kv_block_id[:2] == dummy_block_idx) assert torch.equal( dynamic_context.request_last_kv_block_offset[:2], - torch.tensor([2, 1], device='cuda', dtype=torch.int32), + torch.tensor([2, 1], device='cpu', dtype=torch.int32), ) assert torch.equal( dynamic_context.request_metadata["termination_id"][:2], - torch.tensor([7.0, 8.0], device='cuda'), + torch.tensor([7.0, 8.0], device='cpu'), ) @pytest.mark.internal @@ -569,7 +569,7 @@ def test_add_dummy_requests_parallel_hybrid_allocates_mamba(self): request = DynamicInferenceRequest( request_id=55, - prompt_tokens=torch.arange(0, 5, device='cuda'), + prompt_tokens=torch.arange(0, 5, device='cpu'), sampling_params=SamplingParams(num_tokens_to_generate=4, termination_id=9), ) @@ -577,6 +577,10 @@ def test_add_dummy_requests_parallel_hybrid_allocates_mamba(self): mamba_idx = dynamic_context.mamba_metadata.request_to_mamba_state_idx[0].item() assert mamba_idx >= 0 + + # Mamba state zeroing is deferred until transfer_bookkeeping_to_gpu(). + dynamic_context.initialize_attention_state() + dynamic_context.transfer_bookkeeping_to_gpu() assert torch.all(dynamic_context.mamba_conv_states[:, mamba_idx] == 0) assert torch.all(dynamic_context.mamba_ssm_states[:, mamba_idx] == 0) @@ -597,7 +601,7 @@ def test_add_dummy_requests_parallel_decode_does_not_count_as_prefill(self): request = DynamicInferenceRequest( request_id=5, - prompt_tokens=torch.arange(0, 1, device='cuda'), + prompt_tokens=torch.arange(0, 1, device='cpu'), sampling_params=SamplingParams(num_tokens_to_generate=1, termination_id=2), ) @@ -663,10 +667,10 @@ def test_update_request(self, is_hybrid_model: bool): is_hybrid_model=is_hybrid_model, ) - active_requests_mask = torch.Tensor([1, 0, 1, 1, 1, 0, 0, 1]).cuda().int() - next_tokens = torch.arange(2, 10, device='cuda').int() + active_requests_mask = torch.Tensor([1, 0, 1, 1, 1, 0, 0, 1]).int() + next_tokens = torch.arange(2, 10, device='cpu').int() dynamic_context.paused_request_count = 2 - dynamic_context.paused_tokens = torch.Tensor([0, 1]).cuda().int() + dynamic_context.paused_tokens = torch.Tensor([0, 1]).int() dynamic_context.total_request_count = 5 # Total req count should be equal to paused + num elements in active request mask. @@ -723,7 +727,7 @@ def test_update_request(self, is_hybrid_model: bool): # Then set up the test data dynamic_context.request_ids[0:10] = torch.tensor( - [0, 1, 5, 6, 4, 2, 9, 7, 8, 9], device=torch.cuda.current_device() + [0, 1, 5, 6, 4, 2, 9, 7, 8, 9], device='cpu' ) # Now verify the values @@ -850,12 +854,12 @@ def test_release_memory_blocks_for_finished_requests(self, is_hybrid_model): # Create an active_requests_mask where requests 0, 2, and 4 are finished (0), # and requests 1 and 3 are still active (1) - active_requests_mask = torch.tensor([0, 1, 0, 1, 0], device=torch.cuda.current_device()) + active_requests_mask = torch.tensor([0, 1, 0, 1, 0], device='cpu') # Call update_requests with these parameters dynamic_context.update_requests( active_requests_mask=active_requests_mask, - new_tokens=torch.tensor([10, 11, 12, 13, 14], device=torch.cuda.current_device()), + new_tokens=torch.tensor([10, 11, 12, 13, 14], device='cpu'), ) # After the update, we should have released 3 blocks (for requests 0, 2, and 4) @@ -939,12 +943,12 @@ def test_finished_requests_with_multiple_blocks(self, is_hybrid_model): dynamic_context.mamba_metadata.mamba_state_free_slot_count -= 1 # Create an active_requests_mask where all requests are finished - active_requests_mask = torch.tensor([0, 0, 0], device=torch.cuda.current_device()) + active_requests_mask = torch.tensor([0, 0, 0], device='cpu') # Call update_requests with these parameters dynamic_context.update_requests( active_requests_mask=active_requests_mask, - new_tokens=torch.tensor([10, 11, 12], device=torch.cuda.current_device()), + new_tokens=torch.tensor([10, 11, 12], device='cpu'), ) # After the update, we should have released all 6 blocks and have 0 active requests @@ -994,7 +998,7 @@ def test_mamba_states_cache(self, is_hybrid_model: bool): dynamic_context.add_request( DynamicInferenceRequest( request_id=0, - prompt_tokens=torch.arange(0, context_length, dtype=torch.long, device='cuda'), + prompt_tokens=torch.arange(0, context_length, dtype=torch.long, device='cpu'), sampling_params=SamplingParams( num_tokens_to_generate=dynamic_context.max_tokens - 10 ), @@ -1047,17 +1051,17 @@ def test_calculate_and_store_log_probs(self): # Add a few requests to the context request_data = { 1001: { - "tokens": torch.randint(0, 100, (10,), device='cuda'), + "tokens": torch.randint(0, 100, (10,), device='cpu'), "prefill_len": 10, "initial_token_offset": 0, }, 1002: { - "tokens": torch.randint(0, 100, (5,), device='cuda'), + "tokens": torch.randint(0, 100, (5,), device='cpu'), "prefill_len": 5, "initial_token_offset": 10, }, 1003: { - "tokens": torch.randint(0, 100, (7,), device='cuda'), + "tokens": torch.randint(0, 100, (7,), device='cpu'), "prefill_len": 7, "initial_token_offset": 15, }, @@ -1081,7 +1085,12 @@ def test_calculate_and_store_log_probs(self): # Simulate prefill step total_active_tokens = dynamic_context.active_token_count vocab_size = 50000 - # logits will have shape [1, total_active_tokens, vocab_size] + + # Populate gpu_view for calculate_log_probs (which reads from gpu_view). + dynamic_context.initialize_attention_state() + dynamic_context.transfer_bookkeeping_to_gpu() + + # logits and new_tokens must be on GPU (calculate_log_probs uses gpu_view). prefill_logits = torch.randn( 1, total_active_tokens, vocab_size, device='cuda', dtype=torch.float32 ) @@ -1122,12 +1131,16 @@ def test_calculate_and_store_log_probs(self): # Simulate decode step # All requests are active, so the mask will be all ones for the current active requests - active_requests_mask = torch.ones(dynamic_context.total_request_count, device='cuda').int() + active_requests_mask = torch.ones(dynamic_context.total_request_count, device='cpu').int() dynamic_context.update_requests( active_requests_mask=active_requests_mask, new_tokens=prefill_new_tokens ) + # Populate gpu_view again after update_requests modified bookkeeping state. + dynamic_context.initialize_attention_state() + dynamic_context.transfer_bookkeeping_to_gpu() + # Generate new logits for the decode step. Now each request contributes 1 token. decode_logits = torch.randn( 1, num_active_requests, vocab_size, device='cuda', dtype=torch.float32 @@ -1153,7 +1166,7 @@ def test_calculate_and_store_log_probs(self): # Add a new prefill request to the existing context new_request_id = 1004 - new_request_tokens = torch.randint(0, 100, (12,), device='cuda').long() + new_request_tokens = torch.randint(0, 100, (12,), device='cpu').long() new_request_prefill_len = new_request_tokens.shape[0] initial_token_offset_new_request = dynamic_context.active_token_count dynamic_context.add_request( @@ -1175,6 +1188,7 @@ def test_calculate_and_store_log_probs(self): # This step will involve both prefill (for the new request) and decode (for existing requests). dynamic_context.initialize_attention_state() + dynamic_context.transfer_bookkeeping_to_gpu() total_active_tokens_mixed_step = dynamic_context.active_token_count mixed_step_logits = torch.randn( @@ -1299,7 +1313,7 @@ def test_pipeline_parallel_uneven_layers(self): ), ) - # Collect the total block counts on each rank + # Collect the total block counts on each rank (CUDA needed for NCCL all_gather) local_total_blocks = torch.tensor( [context.kv_block_allocator.total_count], device='cuda', dtype=torch.long ) @@ -1654,14 +1668,14 @@ def test_chunked_prefill_state_preserved_across_decode_completions(self): dynamic_context.add_request( DynamicInferenceRequest( request_id=10, - prompt_tokens=torch.arange(0, 2, device='cuda'), + prompt_tokens=torch.arange(0, 2, device='cpu'), sampling_params=SamplingParams(num_tokens_to_generate=10), ) ) dynamic_context.add_request( DynamicInferenceRequest( request_id=11, - prompt_tokens=torch.arange(0, 2, device='cuda'), + prompt_tokens=torch.arange(0, 2, device='cpu'), sampling_params=SamplingParams(num_tokens_to_generate=10), ) ) @@ -1669,7 +1683,7 @@ def test_chunked_prefill_state_preserved_across_decode_completions(self): # Add Chunk 1 of the chunked prefill request req_999 = DynamicInferenceRequest( request_id=999, - prompt_tokens=torch.arange(0, 8, device='cuda'), + prompt_tokens=torch.arange(0, 8, device='cpu'), sampling_params=SamplingParams(num_tokens_to_generate=10), ) dynamic_context.add_request(req_999, prefill_chunk_length=4) @@ -1683,8 +1697,8 @@ def test_chunked_prefill_state_preserved_across_decode_completions(self): assert kv_block_before != -1 # Step 1: Forward pass for all 3 requests - active_requests_mask = torch.tensor([1, 1, 1], dtype=torch.int32, device='cuda') - new_tokens = torch.tensor([100, 101, 102], dtype=torch.int32, device='cuda') + active_requests_mask = torch.tensor([1, 1, 1], dtype=torch.int32, device='cpu') + new_tokens = torch.tensor([100, 101, 102], dtype=torch.int32, device='cpu') dynamic_context.update_requests(active_requests_mask, new_tokens) # At this point, req 999 is hidden at index 2. total_request_count is 2 (req 10, 11). @@ -1692,8 +1706,8 @@ def test_chunked_prefill_state_preserved_across_decode_completions(self): assert dynamic_context.request_ids[2].item() == 999 # Step 2: Forward pass where req 10 finishes, req 11 continues. Req 999 is NOT scheduled. - active_requests_mask = torch.tensor([0, 1], dtype=torch.int32, device='cuda') - new_tokens = torch.tensor([103, 104], dtype=torch.int32, device='cuda') + active_requests_mask = torch.tensor([0, 1], dtype=torch.int32, device='cpu') + new_tokens = torch.tensor([103, 104], dtype=torch.int32, device='cpu') dynamic_context.update_requests(active_requests_mask, new_tokens) # At this point, req 10 is evicted. Req 11 shifts to index 0. total_request_count becomes 1. @@ -1756,14 +1770,14 @@ def test_chunked_prefill_all_active_requests_finish_while_hidden(self): dynamic_context.add_request( DynamicInferenceRequest( request_id=10, - prompt_tokens=torch.arange(0, 2, device='cuda'), + prompt_tokens=torch.arange(0, 2, device='cpu'), sampling_params=SamplingParams(num_tokens_to_generate=10), ) ) dynamic_context.add_request( DynamicInferenceRequest( request_id=11, - prompt_tokens=torch.arange(0, 2, device='cuda'), + prompt_tokens=torch.arange(0, 2, device='cpu'), sampling_params=SamplingParams(num_tokens_to_generate=10), ) ) @@ -1771,7 +1785,7 @@ def test_chunked_prefill_all_active_requests_finish_while_hidden(self): # Add Chunk 1 of a chunked prefill request req_999 = DynamicInferenceRequest( request_id=999, - prompt_tokens=torch.arange(0, 8, device='cuda'), + prompt_tokens=torch.arange(0, 8, device='cpu'), sampling_params=SamplingParams(num_tokens_to_generate=10), ) dynamic_context.add_request(req_999, prefill_chunk_length=4) @@ -1781,8 +1795,8 @@ def test_chunked_prefill_all_active_requests_finish_while_hidden(self): assert kv_block_before != -1 # Step 1: All 3 requests are active, process forward pass - active_requests_mask = torch.tensor([1, 1, 1], dtype=torch.int32, device='cuda') - new_tokens = torch.tensor([100, 101, 102], dtype=torch.int32, device='cuda') + active_requests_mask = torch.tensor([1, 1, 1], dtype=torch.int32, device='cpu') + new_tokens = torch.tensor([100, 101, 102], dtype=torch.int32, device='cpu') dynamic_context.update_requests(active_requests_mask, new_tokens) # Chunked prefill is now hidden at position 2, total_request_count = 2 @@ -1791,8 +1805,8 @@ def test_chunked_prefill_all_active_requests_finish_while_hidden(self): # Step 2: Both decode requests finish, chunked prefill NOT scheduled this step. # This must NOT crash even though active_request_count becomes 0. - active_requests_mask = torch.tensor([0, 0], dtype=torch.int32, device='cuda') - new_tokens = torch.tensor([103, 104], dtype=torch.int32, device='cuda') + active_requests_mask = torch.tensor([0, 0], dtype=torch.int32, device='cpu') + new_tokens = torch.tensor([103, 104], dtype=torch.int32, device='cpu') dynamic_context.update_requests(active_requests_mask, new_tokens) # total_request_count should be 0 (both finished, chunked prefill hidden) @@ -1839,10 +1853,10 @@ def test_update_requests_speculative(self): ctx.request_to_kv_block_ids[:2, 0] = torch.tensor([0, 1]) ctx.request_last_kv_block_id[:2] = torch.tensor([0, 1]) - active_requests_mask = torch.tensor([1, 1], device='cuda') - new_tokens = torch.tensor([99, 100], device='cuda') # Sampled tokens + active_requests_mask = torch.tensor([1, 1], device='cpu') + new_tokens = torch.tensor([99, 100], device='cpu') # Sampled tokens new_speculative_tokens = torch.tensor( - [[991, 1001], [992, 1002]], device='cuda' + [[991, 1001], [992, 1002]], device='cpu' ) # Spec tokens ctx.update_requests( @@ -1854,15 +1868,14 @@ def test_update_requests_speculative(self): # Each request generates 1 (sampled) + 2 (speculative) = 3 tokens. assert ctx.active_token_count == 6 assert torch.equal( - ctx.request_query_lengths[:2], torch.tensor([3, 3], dtype=torch.int32, device='cuda') + ctx.request_query_lengths[:2], torch.tensor([3, 3], dtype=torch.int32, device='cpu') ) assert torch.equal( - ctx.request_kv_length_offsets[:2], - torch.tensor([6, 9], dtype=torch.int32, device='cuda'), + ctx.request_kv_length_offsets[:2], torch.tensor([6, 9], dtype=torch.int32, device='cpu') ) # Check interleaving: [sampled_1, spec1_1, spec2_1, sampled_2, spec1_2, spec2_2] - expected_tokens = torch.tensor([99, 991, 992, 100, 1001, 1002], device='cuda') + expected_tokens = torch.tensor([99, 991, 992, 100, 1001, 1002], device='cpu') assert torch.equal(ctx.token_to_input_ids[:6], expected_tokens) @pytest.mark.internal @@ -1903,9 +1916,9 @@ def test_speculative_boundary_crossing(self): ctx.request_to_kv_block_ids[0, 0] = first_block ctx.request_last_kv_block_id[0] = first_block - active_requests_mask = torch.tensor([1], device='cuda') - new_tokens = torch.tensor([50], device='cuda') - new_speculative_tokens = torch.tensor([[51], [52]], device='cuda') + active_requests_mask = torch.tensor([1], device='cpu') + new_tokens = torch.tensor([50], device='cpu') + new_speculative_tokens = torch.tensor([[51], [52]], device='cpu') # Run update_requests natively. It will automatically: # 1. Detect the boundary crossing and pause the request. @@ -1929,7 +1942,7 @@ def test_speculative_boundary_crossing(self): # Token 1 (offset 3) -> first_block # Token 2 (offset 4) -> second_block expected_blocks = torch.tensor( - [first_block, first_block, second_block], dtype=torch.int, device='cuda' + [first_block, first_block, second_block], dtype=torch.int, device='cpu' ) assert torch.equal(ctx.token_to_block_idx[:3], expected_blocks) @@ -1979,10 +1992,10 @@ def test_paused_speculative_tokens_tracking(self): ctx.kv_block_allocator.total_avail = 0 ctx.kv_block_allocator.paused_count = 100 # Ensure it doesn't get completely evicted either - active_requests_mask = torch.tensor([1, 1], device='cuda') - new_tokens = torch.tensor([99, 100], device='cuda') # Sampled + active_requests_mask = torch.tensor([1, 1], device='cpu') + new_tokens = torch.tensor([99, 100], device='cpu') # Sampled new_speculative_tokens = torch.tensor( - [[991, 1001], [992, 1002]], device='cuda' + [[991, 1001], [992, 1002]], device='cpu' ) # Speculative # In update_requests, request 0 will be paused to allocate a new block. @@ -2004,7 +2017,7 @@ def test_paused_speculative_tokens_tracking(self): assert ctx.paused_tokens[0].item() == 99 assert torch.equal( - ctx.paused_speculative_tokens[:, 0], torch.tensor([991, 992], device='cuda') + ctx.paused_speculative_tokens[:, 0], torch.tensor([991, 992], device='cpu') ) @pytest.mark.internal @@ -2043,8 +2056,8 @@ def test_swap_book_keeping_tensors_with_speculative_tokens(self): ctx = DynamicInferenceContext(model_config=model_config, inference_config=inference_config) ctx.request_ids[:2] = torch.tensor([10, 11]) - next_tokens = torch.tensor([99, 100], device='cuda') - new_speculative_tokens = torch.tensor([[991, 1001], [992, 1002]], device='cuda') + next_tokens = torch.tensor([99, 100], device='cpu') + new_speculative_tokens = torch.tensor([[991, 1001], [992, 1002]], device='cpu') ctx._swap_book_keeping_tensors( src_idxs=torch.tensor([0]), @@ -2053,10 +2066,10 @@ def test_swap_book_keeping_tensors_with_speculative_tokens(self): new_speculative_tokens=new_speculative_tokens, ) - assert torch.equal(ctx.request_ids[:2], torch.tensor([11, 10], device='cuda')) - assert torch.equal(next_tokens[:2], torch.tensor([100, 99], device='cuda')) + assert torch.equal(ctx.request_ids[:2], torch.tensor([11, 10], device='cpu')) + assert torch.equal(next_tokens[:2], torch.tensor([100, 99], device='cpu')) assert torch.equal( - new_speculative_tokens[:, :2], torch.tensor([[1001, 991], [1002, 992]], device='cuda') + new_speculative_tokens[:, :2], torch.tensor([[1001, 991], [1002, 992]], device='cpu') ) @pytest.mark.internal @@ -2087,9 +2100,9 @@ def test_update_requests_with_finished_requests_and_speculative_tokens(self): ctx.request_last_kv_block_id[:3] = torch.tensor([0, 1, 2]) ctx.request_kv_block_counts[:3] = 1 - active_requests_mask = torch.tensor([1, 0, 1], device='cuda') - new_tokens = torch.tensor([99, 100, 101], device='cuda') - new_speculative_tokens = torch.tensor([[991, 1001, 1011], [992, 1002, 1012]], device='cuda') + active_requests_mask = torch.tensor([1, 0, 1], device='cpu') + new_tokens = torch.tensor([99, 100, 101], device='cpu') + new_speculative_tokens = torch.tensor([[991, 1001, 1011], [992, 1002, 1012]], device='cpu') ctx.update_requests( active_requests_mask=active_requests_mask, @@ -2100,13 +2113,13 @@ def test_update_requests_with_finished_requests_and_speculative_tokens(self): # req1 is finished. req2 moves to req1's position. assert ctx.total_request_count == 2 assert torch.equal( - ctx.request_ids[:2], torch.tensor([10, 12], device='cuda', dtype=torch.int32) + ctx.request_ids[:2], torch.tensor([10, 12], device='cpu', dtype=torch.int32) ) # Check interleaving for req0 and req2 # req0: [99, 991, 992] # req2: [101, 1011, 1012] - expected_tokens = torch.tensor([99, 991, 992, 101, 1011, 1012], device='cuda') + expected_tokens = torch.tensor([99, 991, 992, 101, 1011, 1012], device='cpu') assert torch.equal(ctx.token_to_input_ids[:6], expected_tokens) @pytest.mark.internal @@ -2137,7 +2150,7 @@ def test_chunked_prefill_hidden_state_prevents_token_bloat(self): # 1. Add a standard decode request req_decode = DynamicInferenceRequest( request_id=10, - prompt_tokens=torch.arange(0, 10, device='cuda'), + prompt_tokens=torch.arange(0, 10, device='cpu'), sampling_params=SamplingParams(num_tokens_to_generate=10), ) ctx.add_request(req_decode) @@ -2145,7 +2158,7 @@ def test_chunked_prefill_hidden_state_prevents_token_bloat(self): # 2. Add chunk 1 of a chunked prefill request req_chunked = DynamicInferenceRequest( request_id=42, - prompt_tokens=torch.arange(0, 100, device='cuda'), + prompt_tokens=torch.arange(0, 100, device='cpu'), sampling_params=SamplingParams(num_tokens_to_generate=10), ) ctx.chunked_prefill_request_id = 42 @@ -2155,10 +2168,10 @@ def test_chunked_prefill_hidden_state_prevents_token_bloat(self): assert ctx.active_token_count == 60 # 3. Call update_requests - active_requests_mask = torch.tensor([1, 1], dtype=torch.int32, device='cuda') - new_tokens = torch.tensor([99, 199], dtype=torch.int32, device='cuda') + active_requests_mask = torch.tensor([1, 1], dtype=torch.int32, device='cpu') + new_tokens = torch.tensor([99, 199], dtype=torch.int32, device='cpu') new_spec = torch.tensor( - [[100, 200], [101, 201], [102, 202]], dtype=torch.int32, device='cuda' + [[100, 200], [101, 201], [102, 202]], dtype=torch.int32, device='cpu' ) ctx.update_requests( @@ -2223,13 +2236,13 @@ def test_chunked_prefill_swap_with_speculative_tokens(self): ctx.request_last_kv_block_id[:2] = torch.tensor([0, 1]) ctx.request_kv_block_counts[:2] = 1 - active_requests_mask = torch.tensor([1, 1], device='cuda') + active_requests_mask = torch.tensor([1, 1], device='cpu') # New base tokens: [100 (for prefill), 200 (for decode)] - new_tokens = torch.tensor([100, 200], device='cuda') + new_tokens = torch.tensor([100, 200], device='cpu') # New spec tokens: Col 0 for prefill (dummy), Col 1 for decode (real draft tokens) - new_speculative_tokens = torch.tensor([[101, 201], [102, 202]], device='cuda') + new_speculative_tokens = torch.tensor([[101, 201], [102, 202]], device='cpu') # Trigger update_requests. # It must detect ID 42 is at index 0, and swap it with index 1. @@ -2241,7 +2254,7 @@ def test_chunked_prefill_swap_with_speculative_tokens(self): # 1. Verify the IDs were swapped successfully assert torch.equal( - ctx.request_ids[:2], torch.tensor([99, 42], dtype=torch.int32, device='cuda') + ctx.request_ids[:2], torch.tensor([99, 42], dtype=torch.int32, device='cpu') ) # 2. Verify the Decode request (now at Index 0) correctly flattened its @@ -2249,7 +2262,7 @@ def test_chunked_prefill_swap_with_speculative_tokens(self): # 3. Verify the Prefill request (now at Index 1) is hidden and does NOT # flatten its dummy tokens. expected_flattened_tokens = torch.tensor( - [200, 201, 202], device='cuda' # Decode request (ID 99) + [200, 201, 202], device='cpu' # Decode request (ID 99) ) assert ctx.active_token_count == 3 @@ -2259,7 +2272,7 @@ def test_chunked_prefill_swap_with_speculative_tokens(self): # 4. Verify that the new_speculative_tokens tensor itself was swapped so that # the hidden state perfectly preserves the alignment for subsequent steps. - expected_swapped_spec_tokens = torch.tensor([[201, 101], [202, 102]], device='cuda') + expected_swapped_spec_tokens = torch.tensor([[201, 101], [202, 102]], device='cpu') assert torch.equal( new_speculative_tokens, expected_swapped_spec_tokens ), "new_speculative_tokens was not swapped in-place alongside the request metadata!" @@ -2287,7 +2300,7 @@ def test_speculative_with_prefix_caching_shared_blocks(self): # This avoids the single-token-chunk clamp (effective_prefill >= 2) and # verifies that the prefix skip actually works. tail = 5 - prompt = torch.arange(bs * 3 + tail, device='cuda') + prompt = torch.arange(bs * 3 + tail, device='cpu') # First request registers blocks. req1 = DynamicInferenceRequest( @@ -2349,7 +2362,7 @@ def test_speculative_with_prefix_caching_kv_offset(self): # Use bs * 2 + 5 tokens so the prompt extends past the last full block, # avoiding the single-token-chunk clamp while still testing the skip. tail = 5 - prompt = torch.arange(bs * 2 + tail, device='cuda') + prompt = torch.arange(bs * 2 + tail, device='cpu') # First request. req1 = DynamicInferenceRequest( @@ -2397,7 +2410,7 @@ def test_speculative_update_then_release_with_prefix_caching(self): ctx = DynamicInferenceContext(model_config=model_config, inference_config=inference_config) bs = ctx.block_size_tokens - prompt = torch.arange(bs * 2, device='cuda') + prompt = torch.arange(bs * 2, device='cpu') # Two requests sharing the same prefix. req1 = DynamicInferenceRequest( @@ -2454,7 +2467,7 @@ def test_speculative_boundary_crossing_with_prefix_caching(self): ctx = DynamicInferenceContext(model_config=model_config, inference_config=inference_config) bs = ctx.block_size_tokens - prompt = torch.arange(bs * 2, device='cuda') + prompt = torch.arange(bs * 2, device='cpu') # Request 1: adds prefix blocks. req1 = DynamicInferenceRequest( @@ -2491,9 +2504,9 @@ def test_speculative_boundary_crossing_with_prefix_caching(self): ctx.request_in_prefill_status_tensor[0] = 0 ctx.active_token_count = 2 - active_mask = torch.tensor([1, 1], device='cuda', dtype=torch.int32) - new_tokens = torch.tensor([50, 50], device='cuda') - new_spec = torch.tensor([[51, 51], [52, 52]], device='cuda') + active_mask = torch.tensor([1, 1], device='cpu', dtype=torch.int32) + new_tokens = torch.tensor([50, 50], device='cpu') + new_spec = torch.tensor([[51, 51], [52, 52]], device='cpu') ctx.update_requests( active_requests_mask=active_mask, new_tokens=new_tokens, new_speculative_tokens=new_spec @@ -2535,7 +2548,7 @@ def test_chunked_prefill_prefix_caching_from_hidden_state(self): bs = ctx.block_size_tokens # First request: register prefix blocks (bs * 3 tokens = 3 complete blocks). - first_prompt = torch.arange(bs * 3, device='cuda') + first_prompt = torch.arange(bs * 3, device='cpu') req_first = DynamicInferenceRequest( request_id=1, prompt_tokens=first_prompt.clone(), @@ -2560,9 +2573,9 @@ def test_chunked_prefill_prefix_caching_from_hidden_state(self): ctx.add_request(req2, prefill_chunk_length=bs) # Call update_requests to move req2 to the hidden state - active_requests_mask = torch.tensor([1, 1], dtype=torch.int32, device='cuda') - new_tokens = torch.tensor([99, 199], dtype=torch.int32, device='cuda') - new_spec = torch.tensor([[100, 200], [101, 201]], dtype=torch.int32, device='cuda') + active_requests_mask = torch.tensor([1, 1], dtype=torch.int32, device='cpu') + new_tokens = torch.tensor([99, 199], dtype=torch.int32, device='cpu') + new_spec = torch.tensor([[100, 200], [101, 201]], dtype=torch.int32, device='cpu') ctx.update_requests(active_requests_mask, new_tokens, new_speculative_tokens=new_spec) # Capture active tokens before chunk 2 (which should just be the 3 tokens of req_first) @@ -2604,7 +2617,7 @@ def test_prefix_caching_check_availability_with_speculative(self): ctx = DynamicInferenceContext(model_config=model_config, inference_config=inference_config) bs = ctx.block_size_tokens - prompt = torch.arange(bs * 2, device='cuda') + prompt = torch.arange(bs * 2, device='cpu') # First request registers blocks. req1 = DynamicInferenceRequest( @@ -2654,7 +2667,7 @@ def test_prefix_match_exact_block_boundary(self): bs = ctx.block_size_tokens # req1: 32 tokens (exactly 2 complete blocks) - prompt1 = torch.arange(bs * 2, device='cuda') + prompt1 = torch.arange(bs * 2, device='cpu') req1 = DynamicInferenceRequest( request_id=1, prompt_tokens=prompt1, @@ -2665,7 +2678,7 @@ def test_prefix_match_exact_block_boundary(self): ctx.add_request(req1) # req2: 35 tokens (first 32 tokens match req1) - prompt2 = torch.arange(bs * 2 + 3, device='cuda') + prompt2 = torch.arange(bs * 2 + 3, device='cpu') req2 = DynamicInferenceRequest( request_id=2, prompt_tokens=prompt2, @@ -2712,7 +2725,7 @@ def test_eviction_with_shared_prefix_blocks(self): ctx = DynamicInferenceContext(model_config=model_config, inference_config=inference_config) bs = ctx.block_size_tokens - prompt = torch.arange(bs * 2, device='cuda') + prompt = torch.arange(bs * 2, device='cpu') # Add req1 and req2 with identical prompts req1 = DynamicInferenceRequest( @@ -2752,7 +2765,7 @@ def test_eviction_with_shared_prefix_blocks(self): # Trigger the eviction logic # next_tokens must be sized to total_request_count (1 paused + 1 active = 2) - next_tokens = torch.tensor([50, 51], device='cuda') + next_tokens = torch.tensor([50, 51], device='cpu') evicted_ids = ctx.evict_overflow_paused_requests( active_request_count=1, next_tokens=next_tokens ) @@ -2792,17 +2805,17 @@ def test_oom_during_speculative_boundary_crossing(self): ctx.paused_request_count = 0 ctx.active_token_count = 2 - ctx.request_ids[:2] = torch.tensor([10, 11], device='cuda') + ctx.request_ids[:2] = torch.tensor([10, 11], device='cpu') ctx.request_query_lengths[:2] = 1 ctx.request_kv_block_counts[:2] = 1 # Request 0 offset is 15. Adding 1 sampled + 2 spec = 3 tokens crosses the boundary (16). # Request 1 offset is 5. Adding 3 tokens = 8 (does not cross). ctx.request_kv_length_offsets[:2] = torch.tensor( - [bs - 1, 5], device='cuda', dtype=torch.int32 + [bs - 1, 5], device='cpu', dtype=torch.int32 ) ctx.request_last_kv_block_offset[:2] = torch.tensor( - [bs - 1, 5], device='cuda', dtype=torch.int32 + [bs - 1, 5], device='cpu', dtype=torch.int32 ) blocks = ctx.kv_block_allocator.allocate_memory_blocks(2) @@ -2814,9 +2827,9 @@ def test_oom_during_speculative_boundary_crossing(self): ctx.kv_block_allocator.total_avail = 0 ctx.kv_block_allocator.paused_count = 100 # Prevent immediate eviction out of the system - active_mask = torch.tensor([1, 1], device='cuda', dtype=torch.int32) - new_tokens = torch.tensor([99, 88], device='cuda') - new_spec = torch.tensor([[100, 200], [101, 201]], device='cuda') + active_mask = torch.tensor([1, 1], device='cpu', dtype=torch.int32) + new_tokens = torch.tensor([99, 88], device='cpu') + new_spec = torch.tensor([[100, 200], [101, 201]], device='cpu') # Run update requests ctx.update_requests( @@ -2880,9 +2893,9 @@ def test_speculative_boundary_crossing_at_max_kv_block_count(self): ctx.request_to_kv_block_ids[0, 1] = blocks[1] ctx.request_last_kv_block_id[0] = blocks[1] - active_requests_mask = torch.tensor([1], device='cuda') - new_tokens = torch.tensor([50], device='cuda') - new_speculative_tokens = torch.tensor([[51], [52]], device='cuda') + active_requests_mask = torch.tensor([1], device='cpu') + new_tokens = torch.tensor([50], device='cpu') + new_speculative_tokens = torch.tensor([[51], [52]], device='cpu') # This will pause the request (offset 13 >= 13), then resume it by # allocating a 3rd block at col_idx=2. Without the fix, this raises @@ -2921,7 +2934,7 @@ def test_chunked_prefill_meets_prefix_caching(self): ctx = DynamicInferenceContext(model_config=model_config, inference_config=inference_config) bs = ctx.block_size_tokens - prompt = torch.arange(128, device='cuda') + prompt = torch.arange(128, device='cpu') # Cache req1 (fully processed) req1 = DynamicInferenceRequest( @@ -3153,7 +3166,7 @@ def test_pad_active_slices_no_speculative_tokens(self): num_logits = ctx.num_last_token_logits actual_idxs = ctx.active_logit_idxs[:num_logits] assert torch.equal( - actual_idxs, expected_idxs.to(torch.int32) + actual_idxs, expected_idxs.to(device=actual_idxs.device, dtype=torch.int32) ), f"non-speculative mismatch: {actual_idxs.tolist()} vs {expected_idxs.tolist()}" @pytest.mark.internal diff --git a/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py b/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py index c2deb6d74c7..84898db60d8 100644 --- a/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py +++ b/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py @@ -803,12 +803,12 @@ def test_mamba_intermediate_offsets(self): overall, ) # Penultimate block offset (block 2 boundary) is a valid intermediate - count = msa._intermediate_counts_gpu[1].item() + count = msa._intermediate_counts_cpu[1].item() if count > 0: - offsets = msa._intermediate_offsets_gpu[1, :count].tolist() + offsets = msa._intermediate_offsets_cpu[1, :count].tolist() for o in offsets: assert o > 0 and o % 128 == 0 - assert msa._eos_cache_block_id_gpu[1].item() >= 0 + assert msa._eos_cache_block_id_cpu[1].item() >= 0 # non-aligned prompt produces last_aligned intermediate offset ctx2 = self._mctx(block_size_tokens=bs) @@ -820,12 +820,12 @@ def test_mamba_intermediate_offsets(self): req2b = self._req(ctx2, p2.clone(), request_id=2) req2b._mamba_num_matched_blocks = 2 ctx2.add_request(req2b) - count2 = msa2._intermediate_counts_gpu[1].item() + count2 = msa2._intermediate_counts_cpu[1].item() if count2 > 0: - offsets = msa2._intermediate_offsets_gpu[1, :count2].tolist() + offsets = msa2._intermediate_offsets_cpu[1, :count2].tolist() for o in offsets: assert o > 0 and o % 128 == 0 - assert msa2._eos_cache_block_id_gpu[1].item() < 0 + assert msa2._eos_cache_block_id_cpu[1].item() < 0 # block-aligned prompts set EOS cache block ID ctx3 = self._mctx(block_size_tokens=bs) @@ -834,7 +834,10 @@ def test_mamba_intermediate_offsets(self): req3 = self._req(ctx3, p3.clone(), request_id=2) req3._mamba_num_matched_blocks = 0 ctx3.add_request(req3) - assert ctx3.mamba_slot_allocator._eos_cache_block_id_gpu[1].item() >= 0 + # Deferred Mamba ops execute during transfer. + ctx3.initialize_attention_state() + ctx3.transfer_bookkeeping_to_gpu() + assert ctx3.mamba_slot_allocator._eos_cache_block_id_cpu[1].item() >= 0 # intermediate output buffers are pre-allocated ctx4 = self._mctx() @@ -942,6 +945,7 @@ def test_mixed_batch(self, model_type): # last_token_logits ctx.initialize_attention_state() + ctx.transfer_bookkeeping_to_gpu() logits = torch.randn( 1, ctx.padded_active_token_count, vocab_size, device=torch.cuda.current_device() ) @@ -1067,9 +1071,9 @@ def test_commit_intermediate_states_batched(self): # Set up intermediate offsets: 1 intermediate at src_offset=0 bid0 = ctx.request_to_kv_block_ids[ctx_idx][0].item() - msa._intermediate_block_ids_gpu[ctx_idx, 0] = bid0 - msa._intermediate_offsets_gpu[ctx_idx, 0] = 128 - msa._intermediate_counts_gpu[ctx_idx] = 1 + msa._intermediate_block_ids_cpu[ctx_idx, 0] = bid0 + msa._intermediate_offsets_cpu[ctx_idx, 0] = 128 + msa._intermediate_counts_cpu[ctx_idx] = 1 msa._has_intermediates = True # Set metadata fields that would normally be set by _update_intermediate_offsets @@ -1078,7 +1082,7 @@ def test_commit_intermediate_states_batched(self): # Set up EOS block (block-aligned prompt) eos_bid = ctx.request_to_kv_block_ids[ctx_idx][2].item() - msa._eos_cache_block_id_gpu[ctx_idx] = eos_bid + msa._eos_cache_block_id_cpu[ctx_idx] = eos_bid # Write known patterns to live mamba state for EOS copy mamba_idx = metadata.request_to_mamba_state_idx[ctx_idx].item() diff --git a/tests/unit_tests/inference/engines/test_dynamic_engine.py b/tests/unit_tests/inference/engines/test_dynamic_engine.py index e88924f1af0..24efaea9e1d 100644 --- a/tests/unit_tests/inference/engines/test_dynamic_engine.py +++ b/tests/unit_tests/inference/engines/test_dynamic_engine.py @@ -133,6 +133,7 @@ class DynamicEngineTestConfig: ) force_build_cuda_graphs: bool = False transformer_impl: str = "local" + inference_moe_token_dispatcher_type: str = "nccl" # If False, do not build cuda graphs in the tests, even if # num_cuda_graphs is set. # For tests concerning cuda-graph warmups, we set this to False @@ -330,6 +331,9 @@ def _build_test_env(cls, test_config): inference_sampling_seed=test_config.random_seed, cuda_graph_scope=test_config.cuda_graph_scope, transformer_impl=test_config.transformer_impl, + inference_moe_token_dispatcher_type=( + test_config.inference_moe_token_dispatcher_type + ), normalization=( "RMSNorm" if test_config.transformer_impl == "inference_optimized" @@ -401,6 +405,9 @@ def _build_test_env(cls, test_config): inference_sampling_seed=test_config.random_seed, cuda_graph_scope=test_config.cuda_graph_scope, transformer_impl=test_config.transformer_impl, + inference_moe_token_dispatcher_type=( + test_config.inference_moe_token_dispatcher_type + ), normalization=( "RMSNorm" if test_config.transformer_impl == "inference_optimized" diff --git a/tests/unit_tests/inference/test_mtp_cuda_graph_inference.py b/tests/unit_tests/inference/test_mtp_cuda_graph_inference.py index 60ab7f29a9e..f4e76f6ab7c 100644 --- a/tests/unit_tests/inference/test_mtp_cuda_graph_inference.py +++ b/tests/unit_tests/inference/test_mtp_cuda_graph_inference.py @@ -527,6 +527,9 @@ def _run_mtp(use_cuda_graph): ] ctrl._compute_serial_mtp_and_sample() + # CUDA graph replay is asynchronous, and this test reuses the controller's + # staging buffers immediately for the eager comparison. + torch.cuda.synchronize() return [ ctrl._sampled_mtp_tokens_cuda[d, :active_request_count].clone() diff --git a/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py b/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py index 5846fbd95a0..2c1a9902a1f 100644 --- a/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py +++ b/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py @@ -1038,26 +1038,29 @@ def test_rewind_kv_cache(self, is_hybrid_model): ) self.text_generation_controller.num_speculative_tokens = 3 ctx = self.text_generation_controller.inference_wrapped_model.inference_context + context_device = ctx.request_kv_length_offsets.device ctx.total_request_count = 2 ctx.paused_request_count = 0 - ctx.request_in_prefill_status_tensor = torch.tensor([0, 0], device='cuda') + ctx.request_in_prefill_status_tensor[:2] = torch.tensor( + [0, 0], dtype=torch.int32, device=context_device + ) # Initialize allocator and states ctx.kv_block_allocator.total_avail = 100 - ctx.request_kv_length_offsets[:2] = torch.tensor([10, 15], device='cuda') - ctx.request_kv_block_counts[:2] = torch.tensor([3, 4], device='cuda') + ctx.request_kv_length_offsets[:2] = torch.tensor([10, 15], device=context_device) + ctx.request_kv_block_counts[:2] = torch.tensor([3, 4], device=context_device) # Req 0: offset 2. Rewinding 2 tokens -> offset 0. No block released. # Req 1: offset 1. Rewinding 3 tokens -> offset 2 (prev block). 1 block released. - ctx.request_last_kv_block_offset[:2] = torch.tensor([2, 1], device='cuda') - ctx.request_last_kv_block_id[:2] = torch.tensor([50, 60], device='cuda') + ctx.request_last_kv_block_offset[:2] = torch.tensor([2, 1], device=context_device) + ctx.request_last_kv_block_id[:2] = torch.tensor([50, 60], device=context_device) ctx.request_to_kv_block_ids[:2, :4] = torch.tensor( - [[48, 49, 50, -1], [57, 58, 59, 60]], dtype=torch.int, device='cuda' + [[48, 49, 50, -1], [57, 58, 59, 60]], dtype=torch.int, device=context_device ) if is_hybrid_model: ctx.mamba_metadata.request_to_mamba_state_idx[:2] = torch.tensor( - [0, 1], dtype=torch.int32, device='cuda' + [0, 1], dtype=torch.int32, device=context_device ) ctx.mamba_ssm_states.zero_() ctx.mamba_intermediate_ssm_states.fill_(99) @@ -1076,18 +1079,21 @@ def test_rewind_kv_cache(self, is_hybrid_model): # Assert offsets updated assert torch.equal( ctx.request_last_kv_block_offset[:2], - torch.tensor([0, 2], dtype=torch.int, device='cuda'), + torch.tensor([0, 2], dtype=torch.int, device=context_device), ) assert torch.equal( - ctx.request_kv_length_offsets[:2], torch.tensor([8, 12], dtype=torch.int, device='cuda') + ctx.request_kv_length_offsets[:2], + torch.tensor([8, 12], dtype=torch.int, device=context_device), ) # Assert block counts and IDs updated for boundary crossing assert torch.equal( - ctx.request_kv_block_counts[:2], torch.tensor([3, 3], dtype=torch.int, device='cuda') + ctx.request_kv_block_counts[:2], + torch.tensor([3, 3], dtype=torch.int, device=context_device), ) assert torch.equal( - ctx.request_last_kv_block_id[:2], torch.tensor([50, 59], dtype=torch.int, device='cuda') + ctx.request_last_kv_block_id[:2], + torch.tensor([50, 59], dtype=torch.int, device=context_device), ) # Assert released block is cleared @@ -1275,18 +1281,21 @@ def test_rewind_kv_cache_with_prefix_caching_ref_counts(self): ) ctx = self.text_generation_controller.inference_wrapped_model.inference_context + context_device = ctx.request_kv_length_offsets.device ctx.total_request_count = 2 ctx.paused_request_count = 0 - ctx.request_in_prefill_status_tensor = torch.tensor([0, 0], device='cuda') + ctx.request_in_prefill_status_tensor[:2] = torch.tensor( + [0, 0], dtype=torch.int32, device=context_device + ) # Req 0: 3 blocks, offset 1 in last block. Rewinding 1 token -> no block release. # Req 1: 3 blocks, offset 0 in last block. Rewinding 2 tokens -> crosses back, release block. - ctx.request_kv_length_offsets[:2] = torch.tensor([9, 9], device='cuda') - ctx.request_kv_block_counts[:2] = torch.tensor([3, 3], device='cuda') - ctx.request_last_kv_block_offset[:2] = torch.tensor([1, 0], device='cuda') - ctx.request_last_kv_block_id[:2] = torch.tensor([10, 20], device='cuda') + ctx.request_kv_length_offsets[:2] = torch.tensor([9, 9], device=context_device) + ctx.request_kv_block_counts[:2] = torch.tensor([3, 3], device=context_device) + ctx.request_last_kv_block_offset[:2] = torch.tensor([1, 0], device=context_device) + ctx.request_last_kv_block_id[:2] = torch.tensor([10, 20], device=context_device) ctx.request_to_kv_block_ids[:2, :3] = torch.tensor( - [[8, 9, 10], [18, 19, 20]], dtype=torch.int, device='cuda' + [[8, 9, 10], [18, 19, 20]], dtype=torch.int, device=context_device ) # Set ref counts: block 20 is shared (ref=2), block 10 is exclusive (ref=1). @@ -1321,17 +1330,20 @@ def test_rewind_kv_cache_does_not_release_shared_prefix_blocks(self): ) ctx = self.text_generation_controller.inference_wrapped_model.inference_context + context_device = ctx.request_kv_length_offsets.device ctx.total_request_count = 1 ctx.paused_request_count = 0 - ctx.request_in_prefill_status_tensor = torch.tensor([0], device='cuda') + ctx.request_in_prefill_status_tensor[:1] = torch.tensor( + [0], dtype=torch.int32, device=context_device + ) # 4 blocks. Offset 2 in last block. Rewinding 3 crosses into previous block. - ctx.request_kv_length_offsets[:1] = torch.tensor([14], device='cuda') - ctx.request_kv_block_counts[:1] = torch.tensor([4], device='cuda') - ctx.request_last_kv_block_offset[:1] = torch.tensor([2], device='cuda') - ctx.request_last_kv_block_id[:1] = torch.tensor([40], device='cuda') + ctx.request_kv_length_offsets[:1] = torch.tensor([14], device=context_device) + ctx.request_kv_block_counts[:1] = torch.tensor([4], device=context_device) + ctx.request_last_kv_block_offset[:1] = torch.tensor([2], device=context_device) + ctx.request_last_kv_block_id[:1] = torch.tensor([40], device=context_device) ctx.request_to_kv_block_ids[0, :4] = torch.tensor( - [10, 20, 30, 40], dtype=torch.int, device='cuda' + [10, 20, 30, 40], dtype=torch.int, device=context_device ) # Blocks 10, 20 are shared prefix blocks. Block 30, 40 are exclusive. diff --git a/tools/run_dynamic_text_generation_server.py b/tools/run_dynamic_text_generation_server.py index ed5ab473d63..a33b066e8ea 100644 --- a/tools/run_dynamic_text_generation_server.py +++ b/tools/run_dynamic_text_generation_server.py @@ -10,7 +10,7 @@ start_text_gen_server, stop_text_gen_server, ) -from megatron.core.utils import trace_async_exceptions +from megatron.core.utils import configure_nvtx_profiling, trace_async_exceptions from megatron.inference.utils import add_inference_args, get_dynamic_inference_engine from megatron.post_training.arguments import add_modelopt_args from megatron.training import get_args @@ -83,10 +83,18 @@ async def run_text_generation_server( ) initialize_megatron() + args = get_args() + + # Match training's NVTX gating (training.py only flips this when both + # --profile and --nvtx-ranges are set). Otherwise the engine-side + # nvtx_range_push labels (bookkeeping, Decode, _ep_establish_consensus, + # etc.) are no-ops and the inter-step gap is unattributable in nsys. + if args.profile and args.nvtx_ranges: + configure_nvtx_profiling(True) + # Enable return_log_probs to allow prompt logprobs computation for echo=True requests # This sets materialize_only_last_token_logits=False in the inference context, # which is required for lm-eval compatibility (loglikelihood evaluation tasks) - args = get_args() args.return_log_probs = True engine = get_dynamic_inference_engine()