diff --git a/megatron/core/inference/contexts/attention_context/mamba_metadata.py b/megatron/core/inference/contexts/attention_context/mamba_metadata.py index 3e98f0324e6..3d8c6d4f5b8 100644 --- a/megatron/core/inference/contexts/attention_context/mamba_metadata.py +++ b/megatron/core/inference/contexts/attention_context/mamba_metadata.py @@ -445,8 +445,11 @@ def _update_intermediate_metadata( # Pad unused slots with safe defaults for CUDA graph replay: # - chunk_indices=0: reads from chunk 0 (always exists), output ignored - # - abs_positions=d_conv: conv gather reads tokens [0..d_conv-1], - # which are within bounds and produce a valid but unused state + # - abs_positions=d_conv: conv gather reads tokens [0..d_conv-1]. + # These are within bounds only when the prefill has at least + # d_conv tokens; shorter sequences (e.g. small CUDA-graph warmup + # buckets) would overrun the token axis, so _ssm_prefill clamps + # the gather positions into range. The gathered state is unused. if real_count < max_count: self._intermediate_chunk_indices_buffer[real_count:max_count].fill_(0) self._intermediate_abs_positions_buffer[real_count:max_count].fill_(self.d_conv) @@ -464,7 +467,9 @@ def _update_intermediate_metadata( self.intermediate_abs_positions = self._intermediate_abs_positions_buffer[:max_count] else: # No extraction: fill with safe defaults for CUDA graph warmup - # (same rationale as padding comment above) + # (same rationale as padding comment above; abs_positions=d_conv may + # exceed a sub-d_conv warmup sequence, so _ssm_prefill clamps the + # gather positions into range and the gathered state is unused) self._intermediate_chunk_indices_buffer[:max_count] = 0 self._intermediate_abs_positions_buffer[:max_count] = self.d_conv self.intermediate_count = 0 diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 8e8db5d603d..edb632fb74a 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -278,6 +278,13 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC # Prefix caching hit tracking (accumulated, reset by engine after logging). self.prefix_cache_hits = 0 # requests that matched at least one cached block self.prefix_cache_blocks_matched = 0 # total matched blocks across all requests + # Prefill compute accounting (drained into engine accumulators each step). + # computed = prompt tokens actually run through the model this step; + # skipped = prompt tokens whose prefill was skipped via a prefix-cache hit. + # A high skipped fraction confirms prefix caching is saving prefill compute + # (so any per-step latency growth is attention-over-context, not re-prefill). + self.prefix_cache_prefill_computed_tokens = 0 + self.prefix_cache_prefill_skipped_tokens = 0 # Engine step counter (used for logging, metrics, and event tracking) self.step_count = 0 @@ -2539,12 +2546,23 @@ def reset_tensors(self) -> None: self.token_to_block_idx.fill_(-1) self.token_to_local_position_within_kv_block.fill_(0) - def reset_metadata(self) -> None: + def reset_metadata(self, preserve_prefix_cache: bool = False) -> None: """Reset all bookkeeping state: counters, block allocator, attention/mamba state. This must be called after ``initialize_all_tensors()`` and after any suspend/resume cycle to bring the context back to a clean state. + + Args: + preserve_prefix_cache: When True, keep the KV block allocator's prefix-cache + state (hash index, ref counts, cached blocks) intact. Used by the idle + ``dummy_forward`` path, which only needs to clear the transient one-token + step state -- wiping the allocator there would destroy cross-request prefix + reuse for any subsequent request (the engine idles between requests at low + concurrency, especially with EP > 1). """ + # No cache to preserve when prefix caching is off: fall back to a full + # reset so the disabled path is byte-identical to the original behavior. + preserve_prefix_cache = preserve_prefix_cache and self.enable_prefix_caching # Reset request/token counts. self.total_request_count = 0 @@ -2567,7 +2585,8 @@ def reset_metadata(self) -> None: # Reset attention, mamba, and block allocator state. self.reset_attention_state() self.reset_mamba_state() - self.kv_block_allocator.reset() + if not preserve_prefix_cache: + self.kv_block_allocator.reset() self.request_to_kv_block_ids.fill_(-1) # Reset chunked prefill state @@ -2579,7 +2598,7 @@ def reset_metadata(self) -> None: token_count=0, prefill_req_count=0, decode_req_count=0 ) - def reset(self) -> None: + def reset(self, preserve_prefix_cache: bool = False) -> None: """Reset entire context. This method does: @@ -2590,18 +2609,31 @@ def reset(self) -> None: This method is useful after cuda graph warmup iterations, where the context's memory buffer is referenced by the cuda graph system and cannot be deallocated. + + Args: + preserve_prefix_cache: When True, keep the KV and Mamba prefix-cache + state (hash indices, cached blocks/slots, LRU clock) intact. Used by + the idle ``dummy_forward`` path so an idle step between requests does + not destroy cross-request prefix reuse. """ + # No cache to preserve when prefix caching is off: fall back to a full + # reset so the disabled path is byte-identical to the original behavior. + preserve_prefix_cache = preserve_prefix_cache and self.enable_prefix_caching self.reset_tensors() - self.reset_metadata() + self.reset_metadata(preserve_prefix_cache=preserve_prefix_cache) # Reset lifetime counters (not reset in reset_metadata, which is also # called during suspend/resume where these must persist). - self.step_count = 0 - self.prefix_cache_lru_clock = 0 + if not preserve_prefix_cache: + self.step_count = 0 + self.prefix_cache_lru_clock = 0 - # Reset Mamba cache state - if self.mamba_slot_allocator is not None: - self.mamba_slot_allocator.reset() + # Reset Mamba cache state + if self.mamba_slot_allocator is not None: + self.mamba_slot_allocator.reset() + # When preserving prefix cache (idle dummy_forward), keep step_count + # monotonic so the engine's periodic logging cadence + # (step_count % logging_step_interval) still fires for short requests. def current_input_and_position_ids( self, *, num_warmup_tokens: Optional[int] = None @@ -3072,23 +3104,28 @@ def _register_range(start: int, end: int): else: self._pending_mamba_zeros.append(mamba_idx) - # 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, - current_id, - prefix_skip_tokens, - prefill_chunk_length, - num_matched_blocks, - matched_block_ids, - overall_required_blocks, - ) + # compute_and_store_offsets sets CPU state + GPU staging buffers that + # commit_intermediate_states() consumes after the forward pass. Run it for + # EVERY prefill chunk (not just the first): the last complete block of a + # multi-chunk prompt falls in a continuation chunk, and caching its Mamba + # state is precisely what lets a later turn skip prefill on a hybrid model. + # Mamba slot allocation / state restore above stays first-chunk-only. + if self.is_hybrid_model and self.mamba_slot_allocator is not None: + self.mamba_slot_allocator.compute_and_store_offsets( + req, + current_id, + prefix_skip_tokens, + prefill_chunk_length, + num_matched_blocks, + matched_block_ids, + overall_required_blocks, + ) self.active_token_count += effective_prefill_chunk_length self.lifetime_prefill_token_count += effective_prefill_chunk_length + if self.enable_prefix_caching: + self.prefix_cache_prefill_computed_tokens += effective_prefill_chunk_length + self.prefix_cache_prefill_skipped_tokens += prefix_skip_tokens self.total_request_count += 1 self.num_prefill_requests += 1 diff --git a/megatron/core/inference/contexts/mamba_slot_allocator.py b/megatron/core/inference/contexts/mamba_slot_allocator.py index d7f9c055c60..69f19d442bc 100644 --- a/megatron/core/inference/contexts/mamba_slot_allocator.py +++ b/megatron/core/inference/contexts/mamba_slot_allocator.py @@ -410,24 +410,36 @@ def compute_and_store_offsets( overall_required_blocks: Total blocks needed for this request. """ ctx = self.context + bs = ctx.block_size_tokens prompt_len = len(req.prompt_tokens) - num_kv_matched = num_matched_blocks - kv_div_abs = num_kv_matched * ctx.block_size_tokens - last_aligned_abs = (prompt_len // ctx.block_size_tokens) * ctx.block_size_tokens - seq_len = prefill_chunk_length - skip_tokens # effective prefill length - # Compute relative offsets (relative to prefill start after skip) - kv_div_rel = kv_div_abs - skip_tokens - last_aligned_rel = last_aligned_abs - skip_tokens - penultimate_abs = (overall_required_blocks - 1) * ctx.block_size_tokens - penultimate_rel = penultimate_abs - skip_tokens + # Absolute token position (from the prompt start) where THIS chunk's + # computed tokens begin. The first chunk computes from `skip_tokens` (the + # prefix that was skipped); continuation chunks compute from + # `finished_chunk_token_count` (with skip_tokens == 0). Framing the + # boundary offsets against this chunk start -- rather than assuming the + # first chunk -- lets us extract Mamba state at block boundaries that fall + # in ANY chunk. In particular the last complete block of a multi-chunk + # prompt lives in a continuation chunk; it was previously unreachable, so + # non-block-aligned prompts never cached a usable resume boundary and + # later turns could not skip prefill. + chunk_start = req.finished_chunk_token_count + skip_tokens + seq_len = prefill_chunk_length - skip_tokens # tokens computed this chunk + is_last_chunk = req.finished_chunk_token_count + prefill_chunk_length >= prompt_len + + # Candidate absolute block boundaries at which to cache Mamba state. + kv_div_abs = num_matched_blocks * bs + last_aligned_abs = (prompt_len // bs) * bs # last complete block boundary + penultimate_abs = (overall_required_blocks - 1) * bs # Determine mamba_chunk_size from mamba config (128 is the standard SSM kernel chunk size) mamba_chunk_size = 128 - # Build offset list: include if > 0, < seq_len, and % mamba_chunk_size == 0 + # Keep only boundaries that land inside this chunk's computed tokens and on + # a mamba-chunk boundary (required for mid-sequence state extraction). offsets_set = set() - for offset in [kv_div_rel, last_aligned_rel, penultimate_rel]: + for abs_pos in (kv_div_abs, last_aligned_abs, penultimate_abs): + offset = abs_pos - chunk_start if offset > 0 and offset < seq_len and offset % mamba_chunk_size == 0: offsets_set.add(offset) @@ -436,8 +448,8 @@ def compute_and_store_offsets( # CPU bookkeeping writes (no GPU kernel launches). if count > 0: - 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 + abs_tokens_cpu = torch.tensor([chunk_start + o for o in offsets], dtype=torch.int64) + block_indices_cpu = abs_tokens_cpu // bs - 1 bids_cpu = ctx.request_to_kv_block_ids[current_id][block_indices_cpu] self._intermediate_offsets_cpu[current_id, :count] = torch.tensor( @@ -447,9 +459,13 @@ def compute_and_store_offsets( self._has_intermediates = True 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 + # Block-aligned EOS: when the prompt length is exactly block-aligned, the + # request's live final state IS the last block boundary's state and can be + # cached directly. Only valid on the final chunk (otherwise the live state + # is mid-prompt). Non-block-aligned prompts cache their last complete block + # via the intermediate-extraction path above instead. + if is_last_chunk and last_aligned_abs == prompt_len and prompt_len > 0: + last_block_idx = prompt_len // bs - 1 if last_block_idx >= 0: self._eos_cache_block_id_cpu[current_id] = ctx.request_to_kv_block_ids[current_id][ last_block_idx diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index 01a130656a4..80cc133c6b5 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -352,6 +352,8 @@ def reset(self) -> None: # Prefix caching tracking. self._prefix_cache_hits = 0 self._prefix_cache_blocks_matched = 0 + self._prefill_tokens_computed = 0 + self._prefill_tokens_skipped = 0 self._prefix_coordination_waits = 0 # Coordinator state. @@ -1818,31 +1820,67 @@ def schedule_chunked_prefill(self): if request_can_be_added and kv_cache_available and token_partially_can_be_added: # How many tokens we can admit this step. token_budget = self.context.max_tokens - self.context.active_token_count - max_chunk = min(remaining_len, token_budget) + + # Prefix-cache skip: on a request's first chunk, the tokens covered + # by a cached prefix are reused rather than recomputed, so they do + # NOT consume the compute budget. Extend this chunk's SPAN to cover + # the entire skippable prefix plus up to `token_budget` newly computed + # tokens. Without this the span is capped at the budget, forcing the + # rest of a long cached prefix to be re-prefilled over many chunks + # (latency then scales with prompt length instead of the delta). + # add_request() only computes `effective = span - skip` tokens. + prefix_skip = 0 + if prefix_caching_enabled and not is_continuing_chunked_prefill: + (_, _, _, _, prefix_skip, _) = self.context._compute_prefix_match( + req, remaining_len + ) + prefix_skip = min(prefix_skip, remaining_len - 1) # keep >=1 token to run + + computed_budget = min(remaining_len - prefix_skip, token_budget) # Skip CG gating for the continuation of an in-flight chunked prefill: # the request is already mid-flight, deferring it would deadlock progress. if self._cg_admission_gating_active() and not is_continuing_chunked_prefill: - # Snap chunk size to the largest captured-CG boundary within budget. - # Fall back to eager (max_chunk) if no CG shape covers the budget. - snapped_chunk = self._find_cg_chunk_size(max_chunk) - prefill_chunk_length = snapped_chunk if snapped_chunk is not None else max_chunk + # Snap the COMPUTED chunk size to the largest captured-CG boundary + # within budget (skipped tokens don't affect the CG batch shape). + # Fall back to eager (computed_budget) if no CG shape covers it. + snapped_chunk = self._find_cg_chunk_size(computed_budget) + computed_chunk = snapped_chunk if snapped_chunk is not None else computed_budget req.cg_wait_iters = 0 else: - prefill_chunk_length = max_chunk + computed_chunk = computed_budget + + prefill_chunk_length = prefix_skip + computed_chunk # Flash-attn guard: if this chunk would leave exactly 1 token for the - # final chunk, reduce by 1 (or defer if we only have 1 token of budget). + # final chunk, reduce by 1 (or defer if we only have 1 computed token). # See https://github.com/Dao-AILab/flash-attention/issues/1537 # The -1 is safe after CG snapping: is_applicable_for_batch_dim matches on # cg.token_count >= real.token_count, so the snapped CG still covers token_count-1. if remaining_len - prefill_chunk_length == 1: - if prefill_chunk_length > 1: + if computed_chunk > 1: prefill_chunk_length -= 1 else: can_schedule = False break + # add_request recomputes the skip for this exact chunk and applies a + # ">= 2 computed tokens" clamp. When the chunk would compute fewer than + # 2 tokens (tight budget late in a batched step, or a prompt that is + # all-but-one cached) that clamp shrinks the skip and grows the computed + # count by up to one block, which can exceed the token budget + # (TokenOverflowError). Only then re-derive the exact effective length + # add_request will use and defer on overflow (a later full-budget step + # admits the request). For >= 2 computed tokens add_request computes + # exactly this chunk, which already fits the budget. + if prefix_skip > 0 and (prefill_chunk_length - prefix_skip) < 2: + (_, _, _, _, _, actual_effective) = self.context._compute_prefix_match( + req, prefill_chunk_length + ) + if self.context.active_token_count + actual_effective > self.context.max_tokens: + can_schedule = False + break + # Add hashes to pending set (prefix-caching bookkeeping). if prefix_caching_enabled: for block_hash in req.precomputed_block_hashes: @@ -2078,8 +2116,12 @@ async def async_bookkeep( if self.context.enable_prefix_caching: self._prefix_cache_hits += self.context.prefix_cache_hits self._prefix_cache_blocks_matched += self.context.prefix_cache_blocks_matched + self._prefill_tokens_computed += self.context.prefix_cache_prefill_computed_tokens + self._prefill_tokens_skipped += self.context.prefix_cache_prefill_skipped_tokens self.context.prefix_cache_hits = 0 self.context.prefix_cache_blocks_matched = 0 + self.context.prefix_cache_prefill_computed_tokens = 0 + self.context.prefix_cache_prefill_skipped_tokens = 0 # Log KV cache utilization stats to W&B nvtx_range_push("wandb_logging") @@ -2207,6 +2249,36 @@ async def async_bookkeep( self._prefix_cache_hits, self._prefix_cache_blocks_matched, ) + if self.context.enable_prefix_caching: + # Prefill compute actually saved by prefix caching (cumulative). + # computed = prompt tokens run through the model; skipped = prompt + # tokens whose prefill was reused from cache. If skipped% stays high + # while per-step latency grows, the growth is attention over the + # growing KV context, NOT re-prefilling skipped tokens. + _computed = self._prefill_tokens_computed + _skipped = self._prefill_tokens_skipped + _total = _computed + _skipped + output_str += " ... prefill (cumul): computed %d, skipped %d (%.1f%% skipped)" % ( + _computed, + _skipped, + (100.0 * _skipped / _total) if _total > 0 else 0.0, + ) + # Current cache occupancy (utilization). A Mamba durable-slot count + # near its max indicates the cache is saturating and will start + # LRU-evicting cached prefixes (hybrid models can only skip prefill + # where Mamba state is still cached). + kv_alloc = self.context.kv_block_allocator + output_str += " ... prefix cache util: KV %d/%d blocks cached (%d evictable)" % ( + len(kv_alloc.kv_hash_to_block_id), + kv_alloc.total_count, + int(kv_alloc.get_evictable_block_count()), + ) + msa = self.context.mamba_slot_allocator + if msa is not None: + output_str += ", mamba %d/%d durable slots" % ( + msa.max_slots - msa.free_count, + msa.max_slots, + ) if context_state["is_decode_only"]: output_str = f"\033[94m{output_str}\033[0m" logging.info(output_str) 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 a9c5a6a92f7..d325ee21497 100644 --- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py +++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py @@ -1632,8 +1632,12 @@ def dummy_forward(self): # collectives to avoid a hang. self._dummy_serial_mtp_forward() - # clear the context of any temporary state from the dummy forward - context.reset() + # clear the context of any temporary state from the dummy forward, but + # preserve prefix-cache state: a dummy forward runs when the engine is idle + # (e.g. between requests, or to keep EP collectives alive with EP > 1) and + # must not wipe cached KV/Mamba prefixes, or cross-request prefix reuse would + # be destroyed every time the engine briefly idles. + context.reset(preserve_prefix_cache=True) @torch.inference_mode() def _dummy_serial_mtp_forward(self): diff --git a/megatron/core/ssm/mamba_mixer.py b/megatron/core/ssm/mamba_mixer.py index d2c3d3b1c8d..9ac04a60dd5 100644 --- a/megatron/core/ssm/mamba_mixer.py +++ b/megatron/core/ssm/mamba_mixer.py @@ -1027,6 +1027,16 @@ def _ssm_prefill( intermediate_abs_positions.unsqueeze(1).long() + conv_gather_offsets.unsqueeze(0).long() ) # [n, d_conv] + # Clamp into the valid token range. Padding/warmup slots use the + # safe-default abs_position == d_conv, which yields gather indices + # [0..d_conv-1]; when the prefill sequence is shorter than d_conv + # (e.g. a small CUDA-graph warmup bucket with fewer than d_conv + # tokens), those indices overrun the token axis. Clamping keeps the + # gather in bounds. Real slots are always in range, so this is a + # no-op for them, and padding-slot results are never read (callers + # consult per_request_intermediate_counts). + seq_len = xBC_pre_conv.shape[1] + gather_positions = gather_positions.clamp_(0, seq_len - 1) intermediate_conv = xBC_pre_conv[0, gather_positions, :] # [n, d_conv, conv_dim] intermediate_conv_out[:n].copy_(intermediate_conv.transpose(1, 2)) 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 f808cbedcd9..66d3a963786 100644 --- a/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py +++ b/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py @@ -1367,3 +1367,113 @@ def test_routing_survives_prefix_match_lru(self): assert np.allclose(alloc.get_block_routing(b0), routing_b0) assert alloc.get_block_routing(b1) is not None assert np.allclose(alloc.get_block_routing(b1), routing_b1) + + +class TestPrefixCacheReuse(PrefixCachingTestBase): + """Cross-request prefix reuse on hybrid (Mamba) models: + + - reset(preserve_prefix_cache=True) keeps the cache; a plain reset() clears it. + - Per-context prefill token accounting (computed vs skipped). + - Mamba state is extracted for the last complete block of a multi-chunk prompt. + """ + + @pytest.mark.internal + def test_reset_preserves_prefix_cache_when_requested(self): + # LRU + prefix caching enabled: reset(preserve_prefix_cache=True) keeps the + # KV hash index (so an idle dummy_forward does not wipe cross-request reuse), + # while a plain reset() clears it. + ctx = self._ctx(enable_prefix_caching=True) + bs = ctx.block_size_tokens + ctx.add_request(self._req(ctx, self._prompt(bs * 2))) + cached = dict(ctx.kv_block_allocator.kv_hash_to_block_id) + assert len(cached) == 2 + + ctx.reset(preserve_prefix_cache=True) + assert ctx.kv_block_allocator.kv_hash_to_block_id == cached # preserved + + ctx.reset() # default: full reset + assert len(ctx.kv_block_allocator.kv_hash_to_block_id) == 0 # cleared + + @pytest.mark.internal + def test_reset_disabled_ignores_preserve_flag(self): + # When prefix caching is disabled, preserve_prefix_cache=True still performs + # a full reset: step_count returns to 0. + ctx_off = self._ctx(enable_prefix_caching=False) + ctx_off.step_count = 7 + ctx_off.reset(preserve_prefix_cache=True) + assert ctx_off.step_count == 0 + + # With caching ON, preserve keeps step_count monotonic (for logging cadence). + ctx_on = self._ctx(enable_prefix_caching=True) + ctx_on.step_count = 7 + ctx_on.reset(preserve_prefix_cache=True) + assert ctx_on.step_count == 7 + + @pytest.mark.internal + def test_prefill_computed_and_skipped_counters(self): + # A second request that shares a cached prefix should skip that prefix's + # prefill; the per-context counters must reflect computed vs skipped tokens. + ctx = self._ctx(enable_prefix_caching=True) + bs = ctx.block_size_tokens + + ctx.add_request(self._req(ctx, self._prompt(bs * 4), request_id=1)) + assert ctx.prefix_cache_prefill_skipped_tokens == 0 + assert ctx.prefix_cache_prefill_computed_tokens == bs * 4 + + # request 2 shares the first 4 blocks, adds 2 new blocks + req2 = self._req(ctx, self._prompt(bs * 6), request_id=2) + (matched, _, _, _, prefix_skip, _) = ctx._compute_prefix_match(req2, bs * 6) + assert len(matched) == 4 and prefix_skip == bs * 4 + ctx.add_request(req2) + + assert ctx.prefix_cache_prefill_skipped_tokens == bs * 4 + assert ctx.prefix_cache_prefill_computed_tokens == bs * 6 # 4bs + 2bs + + @pytest.mark.internal + def test_mamba_extraction_covers_last_block_of_continuation_chunk(self): + # For a non-block-aligned, multi-chunk prompt, the last complete block lies + # in a continuation chunk. Extraction offsets are chunk-relative, so that + # boundary's Mamba state is recorded when its chunk is scheduled. + ctx = self._ctx( + mamba_config=self._mamba_config(), + prefix_caching_mamba_gb=0.01, + block_size_tokens=256, + max_sequence_length=4096, + ) # mamba prefix caching enabled + bs = ctx.block_size_tokens + assert bs == 256 + msa = ctx.mamba_slot_allocator + + prompt_len = bs * 3 + 64 # 3 complete blocks + a 64-token remainder + req = self._req(ctx, self._prompt(prompt_len)) + ctx.add_request(req) # populates request_to_kv_block_ids[0] + overall_blocks = ctx.request_kv_block_counts[0].item() + assert overall_blocks == 4 # ceil(832 / 256) + + # Simulate the continuation chunk that covers tokens [2*bs, prompt_len): + # finished=2*bs, no prefix skip, the rest of the prompt as the chunk. + req.finished_chunk_token_count = 2 * bs + cont_chunk = prompt_len - 2 * bs + msa.compute_and_store_offsets( + req, + current_id=0, + skip_tokens=0, + prefill_chunk_length=cont_chunk, + num_matched_blocks=0, + matched_block_ids=[], + overall_required_blocks=overall_blocks, + ) + + # last complete block boundary = 3*bs (768); chunk-relative offset = 768-512=256. + last_aligned_abs = (prompt_len // bs) * bs + expected_offset = last_aligned_abs - 2 * bs + count = msa._intermediate_counts_cpu[0].item() + assert count >= 1 + recorded = msa._intermediate_offsets_cpu[0, :count].tolist() + assert expected_offset in recorded + # the recorded boundary maps to the last complete block (index 2) + idx = recorded.index(expected_offset) + assert ( + msa._intermediate_block_ids_cpu[0, idx].item() + == ctx.request_to_kv_block_ids[0][last_aligned_abs // bs - 1].item() + ) diff --git a/tests/unit_tests/inference/engines/test_prefix_caching_cuda_graphs.py b/tests/unit_tests/inference/engines/test_prefix_caching_cuda_graphs.py index b14747d7070..52e231c1c68 100644 --- a/tests/unit_tests/inference/engines/test_prefix_caching_cuda_graphs.py +++ b/tests/unit_tests/inference/engines/test_prefix_caching_cuda_graphs.py @@ -12,7 +12,6 @@ 2. context.using_cuda_graph_this_step() returned True at expected steps. """ -import os import random import types @@ -44,7 +43,7 @@ from megatron.core.transformer.cuda_graphs import CudaGraphManager, _CudagraphGlobalRecord from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.utils import is_fa_min_version -from tests.unit_tests.test_utilities import Utils +from tests.unit_tests.test_utilities import Utils, clear_nvte_env_vars BLOCK_SIZE = 256 VOCAB_SIZE = 10000 @@ -323,6 +322,11 @@ class TestHybridChunkedPrefillIntermediateState: @classmethod def setup_class(cls): Utils.initialize_model_parallel() + random.seed(123) + torch.manual_seed(123) + model_parallel_cuda_manual_seed( + seed=123, inference_rng_tracker=True, use_cudagraphable_rng=False, force_reset_rng=True + ) @classmethod def teardown_class(cls): @@ -450,16 +454,7 @@ def test_hybrid_chunked_prefill_intermediate_state(self): if not sequence_packing_available: pytest.skip(reason) - # Clear NVTE env vars set by conftest set_env fixture. - os.environ.pop('NVTE_FLASH_ATTN', None) - os.environ.pop('NVTE_FUSED_ATTN', None) - os.environ.pop('NVTE_UNFUSED_ATTN', None) - - random.seed(123) - torch.manual_seed(123) - model_parallel_cuda_manual_seed( - seed=123, inference_rng_tracker=True, use_cudagraphable_rng=False, force_reset_rng=True - ) + clear_nvte_env_vars() # conftest's set_env fixture re-sets these per test model = self._create_hybrid_model() mamba_config = MambaInferenceStateConfig.from_model(model) @@ -541,3 +536,50 @@ def collect_finished(result): f"req {req_id}: baseline {baseline_outputs[req_id]} != " f"test {test_outputs[req_id]}" ) + + @torch.inference_mode() + def test_prefill_shorter_than_conv_window(self): + """A prefill captured into a CUDA graph whose token bucket is smaller than the + Mamba conv window (d_conv) generates correctly. + + Conv-state extraction gathers d_conv positions per slot, and unused slots use + abs_position == d_conv (gather indices up to d_conv-1). The CUDA-graph bucket + list always includes a size-1 (tp_size) graph, so a prompt shorter than d_conv + is captured at a bucket whose token layout is shorter than the gather window. + CUDA graphs (num_cuda_graphs) are required to exercise this capture path. + """ + sequence_packing_available, reason = _check_mamba_sequence_packing_support() + if not sequence_packing_available: + pytest.skip(reason) + + clear_nvte_env_vars() # conftest's set_env fixture re-sets these per test + + model = self._create_hybrid_model(num_cuda_graphs=2) + mamba_config = MambaInferenceStateConfig.from_model(model) + device = torch.cuda.current_device() + + d_conv = mamba_config.conv_states_shape[-1] + if d_conv < 2: + pytest.skip(f"d_conv={d_conv} too small to exercise a sub-window prefill") + + # Prompt shorter than the conv window: its prefill chunk snaps to a CUDA-graph + # bucket < d_conv, so the captured graph's token layout is < d_conv. + engine = self._build_engine( + model, + mamba_config, + enable_prefix_caching=True, + enable_chunked_prefill=True, + num_cuda_graphs=2, + ) + short_prompt = torch.arange(0, d_conv - 1, dtype=torch.int64, device=device) + + engine._add_request(self._make_request(0, short_prompt, enable_pc=True)) + outputs = {} + while engine.has_unfinished_requests(): + result = engine.step_modern() + for record in result["finished_request_records"]: + merged = record.merge() + outputs[merged.request_id] = list(merged.generated_tokens) + + # Generation completes and produces the requested number of tokens. + assert len(outputs[0]) == NUM_TOKENS_TO_GENERATE