From c7d6101e61cba361b4f5a6b3fd21b3896a11b807 Mon Sep 17 00:00:00 2001 From: Shanmugam Ramasamy <111910568+shanmugamr1992@users.noreply.github.com> Date: Wed, 8 Jul 2026 12:43:56 -0700 Subject: [PATCH 1/3] Shadow buffer approach to avoid ima Signed-off-by: Shanmugam Ramasamy <111910568+shanmugamr1992@users.noreply.github.com> --- .../inference/contexts/dynamic_context.py | 43 ++++++++++++++++++- 1 file changed, 42 insertions(+), 1 deletion(-) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 19ff501af91..cb2436683a3 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -1246,6 +1246,20 @@ def initialize_all_tensors(self) -> None: max_mamba_chunks=self._max_mamba_chunks, ) + # Double-buffered pinned shadows that source the async H2D in + # transfer_bookkeeping_to_gpu(). The working buffer above is mutated in + # place every step, so copying directly from it with non_blocking=True + # races the in-flight transfer; we snapshot it into an alternating + # shadow and copy from there instead. See transfer_bookkeeping_to_gpu() + # for the full rationale. CUDA events gate reuse of each shadow. + self._h2d_shadow_bufs = [ + torch.empty(_total_bytes, dtype=torch.uint8, device='cpu', pin_memory=True) + for _ in range(2) + ] + self._h2d_shadow_events = [torch.cuda.Event(), torch.cuda.Event()] + self._h2d_shadow_primed = [False, False] + self._h2d_shadow_idx = 0 + # Cache of (input_ids_view, pos_ids_view) keyed by num_tokens. Instead of slicing and # unsqueezing on every new inference step (constructing new TensorImpls at 30-60 us), # we fix the underlying storage so views are reusable across steps. The number of entries @@ -2457,7 +2471,34 @@ def transfer_bookkeeping_to_gpu(self) -> None: # 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) + # + # Double-buffered async H2D: `self._cpu_bookkeeping_buf` is a pinned host + # buffer mutated in place on the very next step (the staging writes above + # plus `initialize_attention_state()`). A plain `non_blocking=True` copy + # straight from it races: the host overwrites bytes while the async H2D + # is still in flight, so the GPU reads corrupted bookkeeping (token/block + # indices) -> async `CUDA error: an illegal memory access`. The + # CUDA-graph warmup loop captures dimensions in a tight loop, making the + # race fire reliably. + # + # Instead of forcing the copy blocking (which stalls the host on the GPU + # every step), we snapshot the freshly-staged buffer into an alternating + # pinned shadow and async-copy *from the shadow*. The working buffer is + # then immediately free to be re-staged by the next step, while the H2D + # keeps overlapping with subsequent CPU work. A per-shadow CUDA event + # guards reuse: before overwriting a shadow we wait on the event for its + # previous in-flight H2D (a no-op in steady state, since one engine step + # >> the few-us copy; the wait only bites inside the tight warmup loop). + idx = self._h2d_shadow_idx + shadow = self._h2d_shadow_bufs[idx] + if self._h2d_shadow_primed[idx]: + self._h2d_shadow_events[idx].synchronize() + # Host-only memcpy (pinned -> pinned); does not stall on the GPU. + shadow.copy_(self._cpu_bookkeeping_buf) + self.gpu_view._buf.copy_(shadow, non_blocking=True) + self._h2d_shadow_events[idx].record() + self._h2d_shadow_primed[idx] = True + self._h2d_shadow_idx = 1 - idx # MHA metadata GPU views were already bound to state_data in # initialize_attention_state(); the H2D above populates the underlying From 7ea1f3b62a511d1a9fe837090520b744d09db3f1 Mon Sep 17 00:00:00 2001 From: Shanmugam Ramasamy <111910568+shanmugamr1992@users.noreply.github.com> Date: Mon, 13 Jul 2026 16:22:00 -0700 Subject: [PATCH 2/3] fix(inference): make bookkeeping H2D copy blocking to avoid IMA The pinned `_cpu_bookkeeping_buf` is re-staged in place on the next step, so a non_blocking H2D copy races with host writes and can corrupt token/block indices on the GPU. Use a blocking copy instead; the per-step sync cost is negligible relative to the forward pass. Signed-off-by: Shanmugam Ramasamy <111910568+shanmugamr1992@users.noreply.github.com> --- .../inference/contexts/dynamic_context.py | 53 ++++--------------- 1 file changed, 11 insertions(+), 42 deletions(-) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index cb2436683a3..2ed1844fe50 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -1246,20 +1246,6 @@ def initialize_all_tensors(self) -> None: max_mamba_chunks=self._max_mamba_chunks, ) - # Double-buffered pinned shadows that source the async H2D in - # transfer_bookkeeping_to_gpu(). The working buffer above is mutated in - # place every step, so copying directly from it with non_blocking=True - # races the in-flight transfer; we snapshot it into an alternating - # shadow and copy from there instead. See transfer_bookkeeping_to_gpu() - # for the full rationale. CUDA events gate reuse of each shadow. - self._h2d_shadow_bufs = [ - torch.empty(_total_bytes, dtype=torch.uint8, device='cpu', pin_memory=True) - for _ in range(2) - ] - self._h2d_shadow_events = [torch.cuda.Event(), torch.cuda.Event()] - self._h2d_shadow_primed = [False, False] - self._h2d_shadow_idx = 0 - # Cache of (input_ids_view, pos_ids_view) keyed by num_tokens. Instead of slicing and # unsqueezing on every new inference step (constructing new TensorImpls at 30-60 us), # we fix the underlying storage so views are reusable across steps. The number of entries @@ -2471,34 +2457,17 @@ def transfer_bookkeeping_to_gpu(self) -> None: # 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. - # - # Double-buffered async H2D: `self._cpu_bookkeeping_buf` is a pinned host - # buffer mutated in place on the very next step (the staging writes above - # plus `initialize_attention_state()`). A plain `non_blocking=True` copy - # straight from it races: the host overwrites bytes while the async H2D - # is still in flight, so the GPU reads corrupted bookkeeping (token/block - # indices) -> async `CUDA error: an illegal memory access`. The - # CUDA-graph warmup loop captures dimensions in a tight loop, making the - # race fire reliably. - # - # Instead of forcing the copy blocking (which stalls the host on the GPU - # every step), we snapshot the freshly-staged buffer into an alternating - # pinned shadow and async-copy *from the shadow*. The working buffer is - # then immediately free to be re-staged by the next step, while the H2D - # keeps overlapping with subsequent CPU work. A per-shadow CUDA event - # guards reuse: before overwriting a shadow we wait on the event for its - # previous in-flight H2D (a no-op in steady state, since one engine step - # >> the few-us copy; the wait only bites inside the tight warmup loop). - idx = self._h2d_shadow_idx - shadow = self._h2d_shadow_bufs[idx] - if self._h2d_shadow_primed[idx]: - self._h2d_shadow_events[idx].synchronize() - # Host-only memcpy (pinned -> pinned); does not stall on the GPU. - shadow.copy_(self._cpu_bookkeeping_buf) - self.gpu_view._buf.copy_(shadow, non_blocking=True) - self._h2d_shadow_events[idx].record() - self._h2d_shadow_primed[idx] = True - self._h2d_shadow_idx = 1 - idx + # This copy MUST be blocking. `_cpu_bookkeeping_buf` is a pinned host + # buffer that is re-staged in place on the very next step (the staging + # writes above plus `initialize_attention_state()`). A non_blocking copy + # lets the host overwrite those bytes while the async H2D is still in + # flight, so the GPU reads corrupted bookkeeping (token/block indices) + # and dereferences out-of-bounds memory -> async `CUDA error: an illegal + # memory access`. The CUDA-graph warmup loop makes the race fire + # reliably. Blocking costs a per-step host<->device sync, but that is + # negligible relative to the forward pass (benchmarked: no measurable + # generation-throughput difference vs. an async double-buffered copy). + self.gpu_view._buf.copy_(self._cpu_bookkeeping_buf, non_blocking=False) # MHA metadata GPU views were already bound to state_data in # initialize_attention_state(); the H2D above populates the underlying From 75dd58f5c9ec6134d7a21545f4cc8990d3762d42 Mon Sep 17 00:00:00 2001 From: Shanmugam Ramasamy <111910568+shanmugamr1992@users.noreply.github.com> Date: Tue, 14 Jul 2026 11:32:54 -0700 Subject: [PATCH 3/3] docs(inference): clarify blocking bookkeeping H2D in docstring Update transfer_bookkeeping_to_gpu docs to match non_blocking=False and explain the host restage race that requires a blocking copy. Signed-off-by: Shanmugam Ramasamy <111910568+shanmugamr1992@users.noreply.github.com> --- megatron/core/inference/contexts/dynamic_context.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 2ed1844fe50..8d82e745e54 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -2412,13 +2412,15 @@ 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 coalesced H2D from the pinned `_cpu_bookkeeping_buf` uses + ``non_blocking=False``: that buffer is re-staged in place on the next + step, so an async copy can race with host writes and corrupt GPU + bookkeeping (see the inline comment at the copy site). The 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]` + and one contiguous GPU buffer; a single memcpy covers the whole + transfer. 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