CUDA graph support for prefix caching on hybrid models by lmcafee-nvidia · Pull Request #3922 · NVIDIA/Megatron-LM

lmcafee-nvidia · 2026-03-18T13:16:01Z

Summary

Make Mamba state extraction compatible with CUDA graphs during prefix caching
Enable CUDA graphs for chunked prefill on hybrid (transformer + Mamba) models
Add comprehensive correctness tests for prefix caching + CUDA graph interactions

Test plan

test_prefix_caching_cuda_graphs.py passes
Existing test_dynamic_prefix_caching.py tests still pass
test_dynamic_engine.py chunked prefill + CUDA graph tests pass

🤖 Generated with Claude Code

Make intermediate SSM and conv state extraction fully CUDA graph compatible by replacing Python loops and .item() calls with pre-allocated fixed-size buffers and vectorized tensor ops. MambaMetadata precomputes chunk indices and absolute positions during update() (before graph capture/replay). MambaSlotAllocator provides pre-allocated output buffers that _ssm_prefill writes to via GPU ops captured in the graph. commit_intermediate_states() reads from these buffers after the graph completes. This removes the requires_mamba_state_extraction guard that previously forced eager mode for prefill steps needing extraction. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Covers every combination of model type (transformer, hybrid) and batch structure (prefill, decode, mixed) — 6 cases total. Each case runs the same prefix-sharing scenario with and without CUDA graphs, asserting token correctness and verifying actual CUDA graph usage at the expected step types. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Remove two blanket guards that disabled CUDA graphs for all hybrid chunked prefill steps: 1. Guard in match_graph_config() that returned None whenever explicit_chunked_prefill was set and prefill requests were present. Also removes the now-unused explicit_chunked_prefill parameter from match_graph_config() and adjust_batch_dims_for_expert_parallelism(). 2. Guard in is_chunked_prefill_enabled() that disabled chunked prefill during CUDA graph capture for hybrid models. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

…h for continuation chunks Test 1 (test_dynamic_engine.py): parameterized test across model type (gpt/mamba) x chunked prefill (off/on) x CUDA graphs (off/on), verifying output tokens match a baseline config. Uses constrained max_tokens to force actual chunking with multiple active prefill requests. Test 2 (test_prefix_caching_cuda_graphs.py): hybrid model test where one request is mid-chunk (Mamba intermediate state being extracted) while another has Mamba state restored from the prefix cache. Bug fix (dynamic_context.py): the Mamba match count check in `_compute_prefix_match` was incorrectly applied to continuation chunks, where `_mamba_num_matched_blocks` (set over the whole prompt) was compared against the KV match count for only the continuation block range. Added `finished == 0` guard so Mamba match logic only runs for the first chunk. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

copy-pr-bot · 2026-03-18T13:16:05Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

github-actions · 2026-03-18T13:16:13Z

This PR has been automatically converted to draft because all PRs must start as drafts.

When you are ready for review, click Ready for Review to begin the review process. This will:

Add the oncall reviewer (optional reviewer)
Add required review teams based on your changes

See the contribution guide for more details.

…for Mamba dynamic-batching prefill The per-request causal_conv1d_fn loop was a legacy fallback incompatible with CUDA graphs. The Triton causal_conv1d_varlen_fn path is strictly superior and was already force-enabled during graph capture. This removes the fallback entirely along with the use_triton_conv1d config/CLI plumbing. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

lmcafee-nvidia · 2026-03-19T01:19:54Z

/ok to test 44b7f29

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

lmcafee-nvidia · 2026-03-19T01:29:06Z

/ok to test 49c7f97

santhnm2 · 2026-03-19T07:20:30Z

+            for i in range(real_prefill_count):
+                seq_start = cu_seqlens_real[i]
+                seq_len = cu_seqlens_real[i + 1] - seq_start
+                num_chunks = max(1, (seq_len + chunk_size - 1) // chunk_size)
+                first_chunk_idx = cumulative_chunks
+                offsets = intermediate_token_offsets[i]
+                count = 0
+                for offset in offsets:
+                    chunk_idx = first_chunk_idx + (offset // chunk_size) - 1
+                    abs_pos = seq_start + offset
+                    chunk_indices_list.append(chunk_idx)
+                    abs_positions_list.append(abs_pos)
+                    count += 1
+                per_request_counts.append(count)
+                cumulative_chunks += num_chunks


As @kvareddy alluded to, can we vectorize this / directly construct the tensors rather than using lists and loops?

done, we now use vectorized GPU ops here. please take another look

santhnm2 · 2026-03-19T07:28:45Z

+        # Commit intermediate states from output buffers
+        intermediate_count = metadata.intermediate_count
+        if intermediate_count > 0:
+            per_request_counts = metadata.per_request_intermediate_counts
+            ssm_offset = 0
+            req_batch_idx = 0
+            for count in per_request_counts:
+                ctx_idx = prefill_start + req_batch_idx
+                block_ids = self._intermediate_block_ids[ctx_idx]
+
+                if count > 0 and block_ids is not None:
+                    for j in range(count):
+                        bid = block_ids[j]
+                        slot = self.allocate_slot(bid)
+
+                        # Copy states from output buffers for all layers
+                        for layer_idx in range(self.num_mamba_layers):
+                            self.ssm_states[layer_idx, slot].copy_(
+                                self._intermediate_ssm_out[layer_idx, ssm_offset + j]
+                            )
+                            self.conv_states[layer_idx, slot].copy_(
+                                self._intermediate_conv_out[layer_idx, ssm_offset + j]
+                            )
+
+                        # Register in mamba hash map
+                        block_hash = ctx.kv_block_allocator.block_hashes[bid].item()
+                        if block_hash > 0:
+                            self.register_block_hash(bid, block_hash)


Same comment about potential vectorization

this method has been partially vectorized, in the sense that block IDs and hashes are batch-transferred from the GPU via a single .tolist(). the loop isn't easily parallelizable though because it may trigger LRU eviction with side effects like deregistration and free-pool mutation

Would it be possible to do the copy into conv_states and ssm_states as a first step and then do the block hash registration as a second step? And if so, maybe we could more effectively parallelize at least the copy into conv_states and ssm_states (maybe with a Triton kernel?)

this method has been completely rewritten into a pipeline of batch-operation sub-methods. I think the biggest improvement here comes from batching the slot allocation (allocate_slots_batch()), since this potentially results in eviction which is costly. please take another look.

…nd GPU syncs - Replace per-layer state copy loop with all-layers-at-once [:, slot] indexing - Replace per-block invalidate_block loop with batch _invalidate_blocks_batch - Replace conv1d metadata Python loop with torch.repeat_interleave on GPU - Replace while-loop padding with extend([x]*n) bulk allocation - Replace inner while chunk boundary loop with list comprehension + extend Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

…or reduced syncs Replace 3 Python lists (_intermediate_offsets, _intermediate_block_ids, _eos_cache_block_id) with fixed-size GPU tensors. Vectorize block ID lookups via GPU gather, metadata computation via cumsum/masked indexing, and hash fetching via batch .tolist(). Reduces ~24 .item() GPU syncs per step to ~3 bulk transfers for 8 prefill requests with 2 offsets each. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

…ttr naming - Extract hardcoded 3 to MAX_INTERMEDIATE_OFFSETS_PER_REQUEST constant with comment explaining the three candidates (KV divergence, last block-aligned, penultimate block boundaries). - Expand "safe defaults" padding comment to explain why chunk_indices=0 and abs_positions=d_conv are safe for CUDA graph replay. - Remove leading underscore from intermediate_ssm_out, intermediate_conv_out, and conv_gather_offsets since they are accessed from mamba_mixer.py. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

svcnvidia-nemo-ci · 2026-03-24T18:06:50Z