From c581e87760da0e3a155b93c719e4db8ae7c210fa Mon Sep 17 00:00:00 2001 From: Keshav Santhanam Date: Tue, 23 Jun 2026 12:00:02 -0700 Subject: [PATCH 01/11] Fix a bug with prefix caching when number of tokens is less than conv window Signed-off-by: Keshav Santhanam (cherry picked from commit 98eadd841bc2d4f82e7758ab5635654cbbece8ab) --- .../contexts/attention_context/mamba_metadata.py | 11 ++++++++--- megatron/core/ssm/mamba_mixer.py | 10 ++++++++++ 2 files changed, 18 insertions(+), 3 deletions(-) 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/ssm/mamba_mixer.py b/megatron/core/ssm/mamba_mixer.py index 4862da0c81a..a11a0c56192 100644 --- a/megatron/core/ssm/mamba_mixer.py +++ b/megatron/core/ssm/mamba_mixer.py @@ -1028,6 +1028,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)) From dd530732772a9679f3d56687e8897be5f2a6edda Mon Sep 17 00:00:00 2001 From: Keshav Santhanam Date: Thu, 25 Jun 2026 09:20:15 -0700 Subject: [PATCH 02/11] Honor ignore_eos in the dynamic chat completions endpoint The /v1/chat/completions endpoint did not read ignore_eos, so it always terminated generation at the model's EOS even when the request asked to ignore it. This makes forced-length generation (min/exact output length) impossible via the chat API. Mirror the /v1/completions endpoint by setting SamplingParams.termination_id = -1 when ignore_eos is requested. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Keshav Santhanam --- .../dynamic_text_gen_server/endpoints/chat_completions.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/chat_completions.py b/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/chat_completions.py index 96de2e19713..0fccc5a5dae 100644 --- a/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/chat_completions.py +++ b/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/chat_completions.py @@ -593,6 +593,12 @@ async def chat_completions(): max_tokens = req.get("max_completion_tokens", None) or req.get("max_tokens", None) + # When ignore_eos is set, disable EOS-based termination (termination_id=-1) + # so generation runs to num_tokens_to_generate exactly. Mirrors the + # /v1/completions endpoint; without it the chat endpoint stops at the + # model's natural EOS and emits fewer tokens than requested. + ignore_eos = bool(_get_non_none(req, "ignore_eos", False)) + sampling_params = SamplingParams( temperature=temperature, top_k=top_k, @@ -602,6 +608,7 @@ async def chat_completions(): num_tokens_to_generate=(int(max_tokens) if max_tokens is not None else None), skip_prompt_log_probs=skip_prompt_log_probs, add_BOS=add_BOS, + termination_id=-1 if ignore_eos else None, ) except ValueError as e: return Response(f"Invalid sampling parameter: {e}", status=400) From cb3ac23fc47fd87c2906e48788657c495d724553 Mon Sep 17 00:00:00 2001 From: Keshav Santhanam Date: Thu, 25 Jun 2026 09:26:33 -0700 Subject: [PATCH 03/11] Preserve prefix cache across the idle dummy_forward reset dummy_forward() runs a full context.reset() to clear its transient one-token state, but that reset also wipes the KV/Mamba prefix-cache allocator state. The engine runs dummy_forward whenever it idles (between requests, and to keep EP all-to-all collectives alive when EP > 1), so at low concurrency the prefix cache was destroyed on every inter-request gap and never accumulated. Add a preserve_prefix_cache flag to reset()/reset_metadata() that skips the allocator/Mamba-slot reset (and keeps step_count monotonic for logging), and have dummy_forward pass it. The flag is gated on enable_prefix_caching, so the prefix-caching-disabled path is byte-identical to before. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Keshav Santhanam --- .../inference/contexts/dynamic_context.py | 43 +++++++++++++++---- .../text_generation_controller.py | 8 +++- 2 files changed, 40 insertions(+), 11 deletions(-) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 35bcffa82ae..fbeba0f250e 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -2491,12 +2491,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 @@ -2517,7 +2528,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 @@ -2529,7 +2541,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: @@ -2540,18 +2552,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 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 b252e013250..2097059a0a1 100644 --- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py +++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py @@ -1496,8 +1496,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): From 03d2908744c4818cad80d3df8cf8316e917ce2b8 Mon Sep 17 00:00:00 2001 From: Keshav Santhanam Date: Thu, 25 Jun 2026 09:28:19 -0700 Subject: [PATCH 04/11] Cache Mamba state for the last complete block of multi-chunk prompts Hybrid prefill-skip needs a cached Mamba state at the boundary a later turn resumes from. compute_and_store_offsets only ran on a request's first prefill chunk and framed extraction offsets against that chunk, so for prompts longer than one chunk the last complete block (which lives in a continuation chunk) was never extracted, and the live-state EOS path only fired for exactly block-aligned prompts. As a result the durable Mamba cache barely populated and cross-turn prefill-skip almost never triggered. Make the offsets chunk-relative (chunk_start = finished_chunk_token_count + skip_tokens) so a boundary in any chunk is extractable, guard the live-state EOS path to the final chunk, and call compute_and_store_offsets on every prefill chunk (slot allocation/restore stays first-chunk-only). Relies on the existing invariant that continuation chunks carry/restore Mamba state. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Keshav Santhanam --- .../inference/contexts/dynamic_context.py | 30 ++++++------ .../contexts/mamba_slot_allocator.py | 48 ++++++++++++------- 2 files changed, 48 insertions(+), 30 deletions(-) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index fbeba0f250e..868b8facf29 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -3047,20 +3047,22 @@ 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 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 From 1669e19ce43b54c0bf2b8e4f6d8c045c5339250a Mon Sep 17 00:00:00 2001 From: Keshav Santhanam Date: Thu, 25 Jun 2026 09:29:55 -0700 Subject: [PATCH 05/11] Skip the whole cached prefix in one chunk, not one chunk's worth The chunked-prefill scheduler sized a request's chunk span to min(remaining, token_budget), charging skipped (cached) tokens against the per-step compute budget. A long cached prefix could therefore only be skipped one budget-window at a time and the rest was re-prefilled across chunks, so prefill latency scaled with prompt length instead of the uncached delta. Size the first chunk as prefix_skip + min(remaining - prefix_skip, budget) so the entire cached prefix is skipped at once and only the delta is computed (add_request charges the budget against the computed length, not the span). Add an explicit budget re-check: add_request's effective>=2 clamp can shrink the skip and inflate the computed count, so validate the exact effective length and defer the request on overflow (a later full-budget step admits it), preventing a TokenOverflowError that crashed the engine under concurrency. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Keshav Santhanam --- .../core/inference/engines/dynamic_engine.py | 55 ++++++++++++++++--- 1 file changed, 47 insertions(+), 8 deletions(-) diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index 1fe486cf7f9..574cd3c4e2e 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -1757,31 +1757,70 @@ 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: From a4336eb53418922ce5f3573cfa9cbd026f10449b Mon Sep 17 00:00:00 2001 From: Keshav Santhanam Date: Thu, 25 Jun 2026 09:34:43 -0700 Subject: [PATCH 06/11] Log prefix-cache utilization and prefill compute saved Add two cumulative, low-overhead metrics to the periodic engine step log (gated on enable_prefix_caching): - prefill (cumul): computed vs skipped prompt tokens (% skipped) -- shows how much prefill compute prefix caching actually saves, so a rising per-step latency can be attributed to attention over the growing KV context rather than re-prefilling. - prefix cache util: KV blocks cached/total (+evictable) and Mamba durable slots used/max -- surfaces cache occupancy and Mamba-slot saturation. The token counters are accumulated on the context per step and drained into engine accumulators alongside the existing hit counters. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Keshav Santhanam --- .../inference/contexts/dynamic_context.py | 10 ++++++ .../core/inference/engines/dynamic_engine.py | 36 +++++++++++++++++++ 2 files changed, 46 insertions(+) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 868b8facf29..eee7bb3bc67 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -277,6 +277,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 @@ -3066,6 +3073,9 @@ def _register_range(start: int, end: int): 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/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index 574cd3c4e2e..2306c2ab7c1 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -351,6 +351,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. @@ -2056,8 +2058,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") @@ -2185,6 +2191,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) From 99b5218f09d9cbc564dfd0594d24489448ed3c5d Mon Sep 17 00:00:00 2001 From: Keshav Santhanam Date: Thu, 25 Jun 2026 10:01:41 -0700 Subject: [PATCH 07/11] Add unit tests for hybrid prefix-caching fixes Cover the changed behavior in tests/unit_tests/inference/contexts/ test_dynamic_prefix_caching.py (TestPrefixCacheReuseFixes): - reset(preserve_prefix_cache=True) keeps the KV hash index while a plain reset() clears it (the idle dummy_forward path), and the flag is a no-op when prefix caching is disabled. - prefill computed/skipped token counters track a prefix-cache hit correctly. - Mamba state extraction reaches the last complete block of a non-block-aligned multi-chunk prompt via chunk-relative offsets (the boundary lives in a continuation chunk). The whole-chunk-skip scheduler sizing and the TokenOverflow defer guard are exercised end-to-end by the hybrid prefix-caching e2e suite at runtime. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Keshav Santhanam --- .../contexts/test_dynamic_prefix_caching.py | 110 ++++++++++++++++++ 1 file changed, 110 insertions(+) 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() + ) From 09b599434d8fdc95a2534416101e58631431e54b Mon Sep 17 00:00:00 2001 From: Keshav Santhanam Date: Thu, 25 Jun 2026 11:50:05 -0700 Subject: [PATCH 08/11] Add unit test for prefill shorter than the Mamba conv window With CUDA graphs enabled the 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 conv-state gather window (d_conv positions per slot, with unused slots using abs_position == d_conv). The test runs such a prefill end-to-end and asserts it generates the requested number of tokens. Also refactor TestHybridChunkedPrefillIntermediateState to seed once in setup_class and use the shared clear_nvte_env_vars() helper instead of repeating the seed/env boilerplate per test (drops the now-unused os import). Co-Authored-By: Claude Opus 4.8 Signed-off-by: Keshav Santhanam --- .../test_prefix_caching_cuda_graphs.py | 66 +++++++++++++++---- 1 file changed, 54 insertions(+), 12 deletions(-) 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 From 15da24fb1426de6a3ee8b26ea7b11d671af9b295 Mon Sep 17 00:00:00 2001 From: Keshav Santhanam Date: Thu, 25 Jun 2026 13:10:55 -0700 Subject: [PATCH 09/11] Linting Signed-off-by: Keshav Santhanam --- megatron/core/inference/engines/dynamic_engine.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index 2306c2ab7c1..c4946df6c7e 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -1816,10 +1816,7 @@ def schedule_chunked_prefill(self): (_, _, _, _, _, actual_effective) = self.context._compute_prefix_match( req, prefill_chunk_length ) - if ( - self.context.active_token_count + actual_effective - > self.context.max_tokens - ): + if self.context.active_token_count + actual_effective > self.context.max_tokens: can_schedule = False break From 64dd1054ec2b7c5f9d142b183f0c8e97d281ab6e Mon Sep 17 00:00:00 2001 From: Keshav Santhanam Date: Mon, 13 Jul 2026 14:42:57 -0700 Subject: [PATCH 10/11] Relax parent-chained hash guarantee Signed-off-by: Keshav Santhanam --- .../inference/contexts/dynamic_context.py | 27 +++++++----- .../contexts/test_dynamic_prefix_caching.py | 41 +++++++++++++++++++ 2 files changed, 58 insertions(+), 10 deletions(-) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 5d6d8a1cdec..0f1e99c1ff6 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -2860,17 +2860,24 @@ def _find_kv_match_count( hashes = req.precomputed_block_hashes[start_block:end_block] kv_hash_to_block = self.kv_block_allocator.kv_hash_to_block_id - # Find longest KV prefix by iterating block hashes from end. - # Parent-chained hashes guarantee: if hash at position N exists, - # all hashes 0..N also exist. So first match from end = longest prefix. - for i in range(len(hashes) - 1, -1, -1): - if hashes[i] in kv_hash_to_block: - num_matched = i + 1 - matched_blocks = [kv_hash_to_block[hashes[j]] for j in range(num_matched)] - parent_hash = hashes[num_matched - 1] - return matched_blocks, parent_hash + # Return the longest contiguous cached prefix by scanning from the start + # and stopping at the first missing block. Eviction is not parent-aware, + # so a chain can have holes (a shallower block evicted while a deeper one + # stays cached); a deeper match is unusable unless every preceding block + # is also resident, since prefix reuse requires the full contiguous KV + # prefix. + matched_blocks = [] + for h in hashes: + block_id = kv_hash_to_block.get(h) + if block_id is None: + break + matched_blocks.append(block_id) + + if not matched_blocks: + return [], 0 - return [], 0 + parent_hash = hashes[len(matched_blocks) - 1] + return matched_blocks, parent_hash def add_request( self, req: DynamicInferenceRequest, prefill_chunk_length: Optional[int] = None 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 66d3a963786..8d9c35b1b2b 100644 --- a/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py +++ b/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py @@ -378,6 +378,47 @@ def test_block_allocation_with_prefix(self): _, _, kv_available = ctx3.check_availability(self._req(ctx3, p3.clone(), request_id=2)) assert kv_available + @pytest.mark.internal + def test_find_kv_match_count_tolerates_eviction_hole(self): + """A chain with a hole (shallower block evicted, deeper one cached) must + not raise; the match is the longest contiguous prefix from the start. + + Eviction is not parent-aware, so it can drop a parent block while a + child of the same chain stays cached. _find_kv_match_count must return + only the contiguous run before the hole instead of indexing the gap. + """ + ctx = self._ctx() + bs = ctx.block_size_tokens + alloc = ctx.kv_block_allocator + + # Cache a 3-block chain [b0, b1, b2] then release it (blocks stay cached). + prompt = self._prompt(bs * 3) + ctx.add_request(self._req(ctx, prompt.clone())) + b0, b1, b2 = self._block_ids(ctx, 0, 3) + req_hashes = self._req(ctx, prompt.clone(), request_id=99).precomputed_block_hashes + h0, h1, h2 = req_hashes + ctx.release_memory_blocks_from_request_indexes(torch.tensor([0])) + ctx.total_request_count = 0 + assert all(h in alloc.kv_hash_to_block_id for h in (h0, h1, h2)) + + # Punch a hole in the middle of the chain (evict b1 only). + alloc._deregister_blocks(torch.tensor([b1])) + assert h1 not in alloc.kv_hash_to_block_id + assert h0 in alloc.kv_hash_to_block_id and h2 in alloc.kv_hash_to_block_id + + # A new request over the full prompt must match only the contiguous + # prefix before the hole ([b0]) and must not raise KeyError on h1. + req2 = self._req(ctx, prompt.clone(), request_id=2) + assert req2.precomputed_block_hashes == req_hashes + matched, parent_hash = ctx._find_kv_match_count(req2, 0, 3) + assert matched == [b0] + assert parent_hash == h0 + + # Hole at the very first block -> no usable prefix. + alloc._deregister_blocks(torch.tensor([b0])) + matched, parent_hash = ctx._find_kv_match_count(req2, 0, 3) + assert matched == [] and parent_hash == 0 + @pytest.mark.internal def test_ref_count_lru(self): ctx = self._ctx() From f28761fe6b3c0ecc4979d8422885e1807b297343 Mon Sep 17 00:00:00 2001 From: Keshav Santhanam Date: Tue, 14 Jul 2026 09:05:47 -0700 Subject: [PATCH 11/11] Remove _find_kv_match_count and chat_completions.py changes Signed-off-by: Keshav Santhanam --- .../inference/contexts/dynamic_context.py | 27 +++++------- .../endpoints/chat_completions.py | 7 ---- .../contexts/test_dynamic_prefix_caching.py | 41 ------------------- 3 files changed, 10 insertions(+), 65 deletions(-) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 0f1e99c1ff6..5d6d8a1cdec 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -2860,24 +2860,17 @@ def _find_kv_match_count( hashes = req.precomputed_block_hashes[start_block:end_block] kv_hash_to_block = self.kv_block_allocator.kv_hash_to_block_id - # Return the longest contiguous cached prefix by scanning from the start - # and stopping at the first missing block. Eviction is not parent-aware, - # so a chain can have holes (a shallower block evicted while a deeper one - # stays cached); a deeper match is unusable unless every preceding block - # is also resident, since prefix reuse requires the full contiguous KV - # prefix. - matched_blocks = [] - for h in hashes: - block_id = kv_hash_to_block.get(h) - if block_id is None: - break - matched_blocks.append(block_id) - - if not matched_blocks: - return [], 0 + # Find longest KV prefix by iterating block hashes from end. + # Parent-chained hashes guarantee: if hash at position N exists, + # all hashes 0..N also exist. So first match from end = longest prefix. + for i in range(len(hashes) - 1, -1, -1): + if hashes[i] in kv_hash_to_block: + num_matched = i + 1 + matched_blocks = [kv_hash_to_block[hashes[j]] for j in range(num_matched)] + parent_hash = hashes[num_matched - 1] + return matched_blocks, parent_hash - parent_hash = hashes[len(matched_blocks) - 1] - return matched_blocks, parent_hash + return [], 0 def add_request( self, req: DynamicInferenceRequest, prefill_chunk_length: Optional[int] = None diff --git a/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/chat_completions.py b/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/chat_completions.py index 0fccc5a5dae..96de2e19713 100644 --- a/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/chat_completions.py +++ b/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/chat_completions.py @@ -593,12 +593,6 @@ async def chat_completions(): max_tokens = req.get("max_completion_tokens", None) or req.get("max_tokens", None) - # When ignore_eos is set, disable EOS-based termination (termination_id=-1) - # so generation runs to num_tokens_to_generate exactly. Mirrors the - # /v1/completions endpoint; without it the chat endpoint stops at the - # model's natural EOS and emits fewer tokens than requested. - ignore_eos = bool(_get_non_none(req, "ignore_eos", False)) - sampling_params = SamplingParams( temperature=temperature, top_k=top_k, @@ -608,7 +602,6 @@ async def chat_completions(): num_tokens_to_generate=(int(max_tokens) if max_tokens is not None else None), skip_prompt_log_probs=skip_prompt_log_probs, add_BOS=add_BOS, - termination_id=-1 if ignore_eos else None, ) except ValueError as e: return Response(f"Invalid sampling parameter: {e}", status=400) 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 8d9c35b1b2b..66d3a963786 100644 --- a/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py +++ b/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py @@ -378,47 +378,6 @@ def test_block_allocation_with_prefix(self): _, _, kv_available = ctx3.check_availability(self._req(ctx3, p3.clone(), request_id=2)) assert kv_available - @pytest.mark.internal - def test_find_kv_match_count_tolerates_eviction_hole(self): - """A chain with a hole (shallower block evicted, deeper one cached) must - not raise; the match is the longest contiguous prefix from the start. - - Eviction is not parent-aware, so it can drop a parent block while a - child of the same chain stays cached. _find_kv_match_count must return - only the contiguous run before the hole instead of indexing the gap. - """ - ctx = self._ctx() - bs = ctx.block_size_tokens - alloc = ctx.kv_block_allocator - - # Cache a 3-block chain [b0, b1, b2] then release it (blocks stay cached). - prompt = self._prompt(bs * 3) - ctx.add_request(self._req(ctx, prompt.clone())) - b0, b1, b2 = self._block_ids(ctx, 0, 3) - req_hashes = self._req(ctx, prompt.clone(), request_id=99).precomputed_block_hashes - h0, h1, h2 = req_hashes - ctx.release_memory_blocks_from_request_indexes(torch.tensor([0])) - ctx.total_request_count = 0 - assert all(h in alloc.kv_hash_to_block_id for h in (h0, h1, h2)) - - # Punch a hole in the middle of the chain (evict b1 only). - alloc._deregister_blocks(torch.tensor([b1])) - assert h1 not in alloc.kv_hash_to_block_id - assert h0 in alloc.kv_hash_to_block_id and h2 in alloc.kv_hash_to_block_id - - # A new request over the full prompt must match only the contiguous - # prefix before the hole ([b0]) and must not raise KeyError on h1. - req2 = self._req(ctx, prompt.clone(), request_id=2) - assert req2.precomputed_block_hashes == req_hashes - matched, parent_hash = ctx._find_kv_match_count(req2, 0, 3) - assert matched == [b0] - assert parent_hash == h0 - - # Hole at the very first block -> no usable prefix. - alloc._deregister_blocks(torch.tensor([b0])) - matched, parent_hash = ctx._find_kv_match_count(req2, 0, 3) - assert matched == [] and parent_hash == 0 - @pytest.mark.internal def test_ref_count_lru(self): ctx = self._ctx()