diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index e2524a9834a..35bcffa82ae 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -2640,13 +2640,46 @@ def last_token_logits(self, logits: Tensor) -> Tensor: ) return logits.squeeze(0)[self.active_logit_idxs[: self.num_last_token_logits], :] + def _find_mamba_match_count( + self, req: DynamicInferenceRequest, start_block: int, end_block: int + ) -> int: + """Find the farthest cached Mamba state within a chunk-local block range. + + Mamba state restore is only valid for blocks that the current chunk also + assigns from the KV cache. Chunked prefill can schedule a prompt prefix + that is shorter than the farthest cached full-prompt Mamba boundary, so + this helper intentionally uses the same block domain as KV matching. + """ + if self.mamba_slot_allocator is None or not req.precomputed_block_hashes: + return 0 + + end_block = min(end_block, len(req.precomputed_block_hashes)) + if start_block >= end_block: + return 0 + + mamba_map = self.mamba_slot_allocator.hash_to_block_id + hashes = req.precomputed_block_hashes[start_block:end_block] + for i in range(len(hashes) - 1, -1, -1): + if hashes[i] in mamba_map: + return i + 1 + return 0 + def _compute_prefix_match( - self, req: DynamicInferenceRequest, prefill_chunk_length: int + self, + req: DynamicInferenceRequest, + prefill_chunk_length: int, + record_mamba_match: bool = False, ) -> Tuple[list, int, int, int, int, int]: """Compute prefix match results and skip counts for a request chunk. Shared by check_availability (budget checks) and add_request (execution). + Args: + req: Request being scheduled. + prefill_chunk_length: Number of prompt tokens considered in this chunk. + record_mamba_match: If True, store the chunk-local executable Mamba + match count on the request for diagnostics/tests. + Returns: Tuple of (matched_block_ids, num_blocks_from_pool, already_allocated_blocks, overall_required_blocks, @@ -2685,7 +2718,11 @@ def _compute_prefix_match( # Only applies to the first chunk (finished == 0); continuation chunks # already had Mamba state restored during the first chunk. if self.is_hybrid_model and self.mamba_slot_allocator is not None and finished == 0: - num_mamba_matched = getattr(req, '_mamba_num_matched_blocks', 0) + num_mamba_matched = self._find_mamba_match_count( + req, already_allocated_blocks, already_allocated_blocks + num_matched + ) + if record_mamba_match: + req._mamba_num_matched_blocks = num_mamba_matched assert ( num_mamba_matched <= num_matched ), f"Mamba match ({num_mamba_matched}) > KV match ({num_matched})" @@ -2705,6 +2742,8 @@ def _compute_prefix_match( else: prefix_skip_tokens = 0 elif self.is_hybrid_model and finished == 0: + if record_mamba_match: + req._mamba_num_matched_blocks = 0 prefix_skip_tokens = 0 # Clamp so that effective_prefill_chunk_length >= 2 when possible. @@ -2831,7 +2870,7 @@ def add_request( overall_required_blocks, prefix_skip_tokens, effective_prefill_chunk_length, - ) = self._compute_prefix_match(req, prefill_chunk_length) + ) = self._compute_prefix_match(req, prefill_chunk_length, record_mamba_match=True) num_matched_blocks = len(matched_block_ids) effective_kv_offset = req.finished_chunk_token_count + prefix_skip_tokens diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index 9d741cae0e9..1fe486cf7f9 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -1544,22 +1544,6 @@ def get_prefix_coordination_metrics(self) -> dict: """ return {"waits": self._prefix_coordination_waits} - def _find_mamba_match_count(self, req: DynamicInferenceRequest) -> int: - """Find farthest block with cached Mamba state by iterating from the end. - - Not all blocks have Mamba state cached in mamba_hash_to_block_id, - only divergence and last-aligned blocks do. Iterating from the end - finds the farthest block with cached state, which is the only one - needed for restore since Mamba state is cumulative. - """ - if not req.precomputed_block_hashes: - return 0 - mamba_map = self.context.mamba_slot_allocator.hash_to_block_id - for i in range(len(req.precomputed_block_hashes) - 1, -1, -1): - if req.precomputed_block_hashes[i] in mamba_map: - return i + 1 - return 0 - def schedule_waiting_requests(self): """Tries to schedule any requests in the waiting pool.""" # Keep track of which requests get scheduled. @@ -1582,11 +1566,6 @@ def schedule_non_chunked_prefill(self): Perform the same original scheduling logic for non-chunked runs """ prefix_caching_enabled = self.context.enable_prefix_caching - mamba_caching_enabled = ( - prefix_caching_enabled - and self.context.is_hybrid_model - and self.context.mamba_slot_allocator is not None - ) if prefix_caching_enabled: pending_block_hashes = set() pending_request_ids = [] @@ -1605,10 +1584,6 @@ def schedule_non_chunked_prefill(self): pending_request_ids.append(self.waiting_request_ids.popleft()) continue - # Find Mamba prefix match before check_availability (sets skip count) - if mamba_caching_enabled: - req._mamba_num_matched_blocks = self._find_mamba_match_count(req) - request_can_be_added, request_tokens_can_be_added, kv_cache_available = ( self.context.check_availability(req) ) @@ -1746,11 +1721,6 @@ def schedule_chunked_prefill(self): - For each request, remaining_prompt_tokens holds the **unprefilled** prompt tokens """ prefix_caching_enabled = self.context.enable_prefix_caching - mamba_caching_enabled = ( - prefix_caching_enabled - and self.context.is_hybrid_model - and self.context.mamba_slot_allocator is not None - ) if prefix_caching_enabled: pending_block_hashes = set() pending_request_ids = [] @@ -1778,10 +1748,6 @@ def schedule_chunked_prefill(self): ) continue - # Find Mamba prefix match for non-continuing requests - if mamba_caching_enabled and not is_continuing_chunked_prefill: - req._mamba_num_matched_blocks = self._find_mamba_match_count(req) - # Use remaining prompt tokens for scheduling decisions remaining_len = len(req.remaining_prompt_tokens) token_partially_can_be_added = self.context.active_token_count < self.context.max_tokens 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 75feb091086..f808cbedcd9 100644 --- a/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py +++ b/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py @@ -689,17 +689,13 @@ def test_mamba_cache_lifecycle(self): ctx6.add_request(self._req(ctx6, p6.clone())) msa6 = ctx6.mamba_slot_allocator self._mamba_allocate_and_register(ctx6, self._block_ids(ctx6, 0, 4)[:2]) - engine6 = _StubEngine(ctx6) - assert engine6._find_mamba_match_count(self._req(ctx6, p6.clone(), request_id=2)) == 2 + req6 = self._req(ctx6, p6.clone(), request_id=2) + assert ctx6._find_mamba_match_count(req6, 0, len(req6.precomputed_block_hashes)) == 2 # no match when no mamba hashes registered ctx7 = self._mctx() ctx7.add_request(self._req(ctx7, self._prompt(bs * 3))) - assert ( - _StubEngine(ctx7)._find_mamba_match_count( - self._req(ctx7, self._prompt(bs * 3), request_id=2) - ) - == 0 - ) + req7 = self._req(ctx7, self._prompt(bs * 3), request_id=2) + assert ctx7._find_mamba_match_count(req7, 0, len(req7.precomputed_block_hashes)) == 0 # allocate, free, re-allocate ctx8 = self._mctx() diff --git a/tests/unit_tests/inference/engines/test_dynamic_engine.py b/tests/unit_tests/inference/engines/test_dynamic_engine.py index ec008d42836..89242ee3182 100644 --- a/tests/unit_tests/inference/engines/test_dynamic_engine.py +++ b/tests/unit_tests/inference/engines/test_dynamic_engine.py @@ -1958,6 +1958,55 @@ def test_prefix_caching_avoid_single_token_effective_chunk(self): f"but got {ctx.request_query_lengths[req_b_idx].item()}." ) + @pytest.mark.internal + @torch.inference_mode() + def test_mamba_match_is_chunk_local_when_chunked_prefill_limits_kv_match(self): + """Mamba restore depth is bounded to KV blocks assigned for the current chunk.""" + block_size = 256 + block_hashes = [111, 222] + req = DynamicInferenceRequest( + request_id=1, + prompt_tokens=torch.arange(512, dtype=torch.int64, device='cuda'), + sampling_params=SamplingParams(num_tokens_to_generate=1), + block_size_tokens=block_size, + enable_prefix_caching=True, + precomputed_block_hashes=block_hashes, + ) + + # This simulates the old scheduler-side full-prompt Mamba match. The + # context must ignore it and record the chunk-local executable count. + req._mamba_num_matched_blocks = 2 + + ctx = DynamicInferenceContext.__new__(DynamicInferenceContext) + ctx.block_size_tokens = block_size + ctx.enable_prefix_caching = True + ctx.is_hybrid_model = True + ctx.kv_block_allocator = types.SimpleNamespace( + kv_hash_to_block_id={block_hashes[0]: 7, block_hashes[1]: 8} + ) + ctx.mamba_slot_allocator = types.SimpleNamespace( + hash_to_block_id={block_hashes[0]: 7, block_hashes[1]: 8} + ) + + ( + matched_block_ids, + num_blocks_from_pool, + already_allocated_blocks, + overall_required_blocks, + prefix_skip_tokens, + effective_prefill_chunk_length, + ) = DynamicInferenceContext._compute_prefix_match( + ctx, req, prefill_chunk_length=211, record_mamba_match=True + ) + + assert matched_block_ids == [7] + assert num_blocks_from_pool == 0 + assert already_allocated_blocks == 0 + assert overall_required_blocks == 1 + assert req._mamba_num_matched_blocks == 1 + assert prefix_skip_tokens == 0 + assert effective_prefill_chunk_length == 211 + @pytest.mark.internal @pytest.mark.skipif( not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching"