diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py index d49ab084a128..a9f84ac23e40 100644 --- a/tensorrt_llm/_torch/attention_backend/flashinfer.py +++ b/tensorrt_llm/_torch/attention_backend/flashinfer.py @@ -702,12 +702,13 @@ def _sync_draft_view(self, target: "FlashInferAttentionMetadata") -> None: destination[:total_blocks].copy_(source[:total_blocks], non_blocking=True) self._vswa_pool_indices_cache[pool_id] = destination - primary_pool_id = self._vswa_layer_to_pool.get(0, 0) - self._vswa_active_pool_id = primary_pool_id + default_pool_id = target._vswa_layer_to_pool.get( + target._default_kv_layer_idx, 0) + self._vswa_active_pool_id = default_pool_id self._host_paged_kv_indices = self._host_pool_indices[ - primary_pool_id] + default_pool_id] self._paged_kv_indices[:total_blocks].copy_( - self._vswa_pool_indices_cache[primary_pool_id][:total_blocks], + self._vswa_pool_indices_cache[default_pool_id][:total_blocks], non_blocking=True) else: self._vswa_pool_indices_cache = None @@ -769,16 +770,17 @@ def update_shared_kv_draft_lengths( def _prepare_full_draft_page_table(self) -> None: """Expose every allocated draft page and use device KV lengths.""" - if self._uses_full_draft_page_table: + if (self._uses_full_draft_page_table + or self._default_kv_layer_idx is None): return assert self.request_ids is not None block_ids_per_seq = self.kv_cache_manager.get_batch_cache_indices( - self.request_ids) + self.request_ids, layer_idx=self._default_kv_layer_idx) self.num_blocks = [len(block_ids) for block_ids in block_ids_per_seq] self.num_context_blocks = 0 self.num_generation_blocks = sum(self.num_blocks) - paged_kv_indices = self.kv_cache_manager.get_batch_cache_indices_flat( - self.request_ids, self.num_blocks) + paged_kv_indices = self._get_default_kv_indices(self.request_ids, + self.num_blocks) self._paged_kv_indices[:paged_kv_indices.numel()].copy_( paged_kv_indices, non_blocking=True) self._host_paged_kv_indices = paged_kv_indices @@ -836,9 +838,35 @@ def __post_init__(self) -> None: super().__post_init__() self._post_init_with_buffers(self.cuda_graph_buffers) + def _get_pp_attention_layer_ids(self) -> list[int]: + """Return global IDs of attention layers assigned to this PP rank.""" + assert self.kv_cache_manager is not None + layer_offsets = getattr(self.kv_cache_manager, "layer_offsets", None) + num_kv_heads = getattr(self.kv_cache_manager, "num_kv_heads_per_layer", + None) + if layer_offsets is None or num_kv_heads is None: + return [0] + return [ + layer_idx for layer_idx, layer_offset in layer_offsets.items() + if num_kv_heads[layer_offset] > 0 + ] + + def _get_default_kv_indices(self, request_ids: list[int], + num_blocks: list[int]) -> torch.Tensor: + """Return page indices from the default attention KV pool.""" + assert self.kv_cache_manager is not None + assert self._default_kv_layer_idx is not None + return self.kv_cache_manager.get_batch_cache_indices_flat( + request_ids, + num_blocks, + layer_idx=self._default_kv_layer_idx, + ) + def _post_init_with_buffers(self, buffers) -> None: capture_graph = self.is_cuda_graph + self._default_kv_layer_idx: Optional[int] = None + if self.workspace_buffer is None: # Note: even though flashinfer only recommends 128 MB, we have to push it # a bit higher to cover all possible CUDA graph cases. If it's too small, @@ -919,25 +947,34 @@ def _post_init_with_buffers(self, buffers) -> None: self._vswa_pool_indices_cache: Optional[Dict[int, torch.Tensor]] = None if self.kv_cache_manager is not None: - blocks_in_primary_pool = self.kv_cache_manager.blocks_in_primary_pool + pp_attention_layer_ids = self._get_pp_attention_layer_ids() + self._default_kv_layer_idx = next(iter(pp_attention_layer_ids), + None) + if self._default_kv_layer_idx is None: + logger.info_once( + "FlashInfer: no attention-backed KV layer on this rank; " + "paged KV metadata disabled", + key="flashinfer_no_attention_backed_kv_layer", + ) + max_num_blocks = 0 + get_buffers = getattr(self.kv_cache_manager, "get_buffers", None) + if get_buffers is not None: + for layer_idx in pp_attention_layer_ids: + layer_buffer = get_buffers(layer_idx) + if layer_buffer is not None: + max_num_blocks = max(max_num_blocks, + layer_buffer.shape[0]) + if pp_attention_layer_ids and max_num_blocks == 0: + max_num_blocks = int( + getattr(self.kv_cache_manager, "blocks_in_primary_pool", 0)) self._paged_kv_indices = self.get_empty( buffers, - (blocks_in_primary_pool, ), + (max_num_blocks, ), dtype=torch.int, cache_name="_paged_kv_indices", capture_graph=capture_graph, ) - # Maximum block count across ALL pools: sizes the VSWA pool - # buffers below. Computed for every model, not just VSWA — - # non-VSWA managers have a single pool, so this stays - # blocks_in_primary_pool for them. - max_num_blocks = blocks_in_primary_pool - if hasattr(self.kv_cache_manager, 'layer_offsets'): - for lid in self.kv_cache_manager.layer_offsets: - lbuf = self.kv_cache_manager.get_buffers(lid) - if lbuf is not None: - max_num_blocks = max(max_num_blocks, lbuf.shape[0]) self._max_num_blocks_per_seq = self.kv_cache_manager.max_blocks_per_seq # Layers may share one page-index list only when they are in the @@ -951,7 +988,7 @@ def _post_init_with_buffers(self, buffers) -> None: layer_space: Dict[int, int] = {} if hasattr(mgr, 'layer_to_pool_mapping_dict') and get_scale: space_ids = {} - for layer_idx in getattr(mgr, 'layer_offsets', {}): + for layer_idx in pp_attention_layer_ids: layer_offset = mgr.layer_offsets[layer_idx] key = (mgr.layer_to_pool_mapping_dict[layer_offset], get_scale(layer_idx)) @@ -1401,6 +1438,20 @@ def _to_int32_tensor(arr: np.ndarray) -> torch.Tensor: assert self.request_ids is not None + if self._default_kv_layer_idx is None: + self.num_blocks = [0] * self.num_seqs + self.num_context_blocks = 0 + self.num_generation_blocks = 0 + self.num_ctx_cached_tokens = 0 + self._cached_token_lens[:self.num_seqs].zero_() + self.paged_kv_indptr_decode[:self.num_seqs + 1].zero_() + self.paged_kv_indptr_prefill[:self.num_seqs + 1].zero_() + self.paged_kv_indptr = self.paged_kv_indptr_decode[:self.num_seqs + + 1] + self._paged_kv_last_page_len[:self.num_seqs].zero_() + self._clean_cached_plans(defer_plan=False) + return + # start and end indices of each sequence in the ragged key and value # for self attention it's the same as qo_indptr so avoid computing twice. if self.is_cross: @@ -1419,8 +1470,8 @@ def _to_int32_tensor(arr: np.ndarray) -> torch.Tensor: self.num_generation_blocks = sum(self.num_blocks[self.num_contexts:]) # indices of used cache blocks for each sequence - paged_kv_indices = self.kv_cache_manager.get_batch_cache_indices_flat( - self.request_ids, self.num_blocks) + paged_kv_indices = self._get_default_kv_indices(self.request_ids, + self.num_blocks) self._paged_kv_indices[:paged_kv_indices.size(0)].copy_( paged_kv_indices, non_blocking=True) @@ -1437,20 +1488,20 @@ def _to_int32_tensor(arr: np.ndarray) -> torch.Tensor: # capturable). if self._vswa_layer_to_pool is not None: unique_pools = set(self._vswa_layer_to_pool.values()) - primary_pool_id = self._vswa_layer_to_pool.get(0, 0) + default_pool_id = self._vswa_layer_to_pool.get( + self._default_kv_layer_idx, 0) # Use dedicated pre-allocated buffers for each pool's indices. # These buffers are created in __post_init__ so their addresses # stay stable across CUDA-graph replays. total_idx = paged_kv_indices.size(0) - primary_buf = getattr(self, f'_vswa_pool_buf_{primary_pool_id}') - primary_buf[:total_idx].copy_(self._paged_kv_indices[:total_idx], - non_blocking=True) + default_buf = getattr(self, f'_vswa_pool_buf_{default_pool_id}') + default_buf[:total_idx].copy_(paged_kv_indices, non_blocking=True) self._vswa_pool_indices_cache = { - primary_pool_id: primary_buf, + default_pool_id: default_buf, } - self._host_pool_indices = {primary_pool_id: paged_kv_indices} + self._host_pool_indices = {default_pool_id: paged_kv_indices} for pool_id in unique_pools: - if pool_id == primary_pool_id: + if pool_id == default_pool_id: continue rep_layer = self._vswa_pool_to_rep_layer[pool_id] pool_indices = \ @@ -1461,7 +1512,7 @@ def _to_int32_tensor(arr: np.ndarray) -> torch.Tensor: non_blocking=True) self._vswa_pool_indices_cache[pool_id] = buf self._host_pool_indices[pool_id] = pool_indices - self._vswa_active_pool_id = primary_pool_id + self._vswa_active_pool_id = default_pool_id # number of tokens in the last cache block used by each sequence, # derived on the host so no GPU arithmetic or sync is needed. @@ -1571,16 +1622,17 @@ def _to_int32_tensor(arr: np.ndarray) -> torch.Tensor: self._do_plan_mla_decode(self._mla_decode_plan_params) self._mla_decode_planned = True - # VSWA: restore primary pool indices as the default. + # VSWA: restore the default pool indices. if (self._vswa_layer_to_pool is not None and self._vswa_pool_indices_cache is not None): - primary_pool_id = self._vswa_layer_to_pool.get(0, 0) + default_pool_id = self._vswa_layer_to_pool.get( + self._default_kv_layer_idx, 0) total_blocks = self.num_generation_blocks + self.num_context_blocks - src = self._vswa_pool_indices_cache[primary_pool_id][:total_blocks] + src = self._vswa_pool_indices_cache[default_pool_id][:total_blocks] self._paged_kv_indices[:total_blocks].copy_(src, non_blocking=True) - self._vswa_active_pool_id = primary_pool_id + self._vswa_active_pool_id = default_pool_id self._host_paged_kv_indices = \ - self._host_pool_indices[primary_pool_id] + self._host_pool_indices[default_pool_id] # CUDA graph + trtllm-gen: update _block_tables and _kv_lens_buffer # so the trtllm-gen decode kernel uses current page indices. diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 2ea7d89a6fbc..0faaefaca1da 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -461,9 +461,6 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int], self.is_linear_attention = linear_attention_metadata is not None # Calculate kv cache blocks for each window size - # FIXME: flashinfer.py accesses kv_cache_manager.blocks_in_primary_pool - # This dependency should be adjusted as it only covers the single window - # case and not VSWA scheme. if is_estimating_kv_cache: # If this is an estimation dry run, we have already calculated the # max_tokens under _util.py::try_prepare_estimation diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index a65e2c70ef3e..8128c509e9c6 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -6172,6 +6172,20 @@ def test_bf16(self): task.evaluate(llm, extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS) + @skip_pre_hopper + def test_bf16_flashinfer(self) -> None: + model_path = f"{llm_models_root()}/Qwen3.5-4B" + with LLM(model_path, + trust_remote_code=True, + max_seq_len=4096, + max_batch_size=32, + kv_cache_config=self.kv_cache_config, + cuda_graph_config=self.cuda_graph_config, + attn_backend="FLASHINFER") as llm: + task = GSM8K(self.MODEL_NAME) + task.evaluate(llm, + extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS) + @skip_pre_hopper def test_fp8(self): model_path = f"{llm_models_root()}/Qwen3.5-4B-FP8" diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index 72010fb83331..d5f68a77674c 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -768,6 +768,7 @@ accuracy/test_llm_api_pytorch.py::TestQwen3_5_397B_A17B::test_nvfp4[tep4_cutedsl accuracy/test_llm_api_pytorch.py::TestQwen3_5_397B_A17B::test_nvfp4[tep4_trtllm] accuracy/test_llm_api_pytorch.py::TestQwen3_5_397B_A17B::test_nvfp4_mtp3_gdn_replay_tep4 accuracy/test_llm_api_pytorch.py::TestQwen3_5_4B::test_bf16 +accuracy/test_llm_api_pytorch.py::TestQwen3_5_4B::test_bf16_flashinfer accuracy/test_llm_api_pytorch.py::TestQwen3_5_4B::test_dflash accuracy/test_llm_api_pytorch.py::TestQwen3_5_4B::test_fp8 accuracy/test_llm_api_pytorch.py::TestQwen3_5_4B::test_fp8_piecewise_cuda_graph diff --git a/tests/integration/test_lists/qa/llm_function_rtx6k.txt b/tests/integration/test_lists/qa/llm_function_rtx6k.txt index a354d3a9a1a6..ba94a0dd885d 100644 --- a/tests/integration/test_lists/qa/llm_function_rtx6k.txt +++ b/tests/integration/test_lists/qa/llm_function_rtx6k.txt @@ -176,6 +176,7 @@ accuracy/test_llm_api_pytorch.py::TestQwen3_5_397B_A17B::test_nvfp4[adp4_trtllm] accuracy/test_llm_api_pytorch.py::TestQwen3_5_397B_A17B::test_nvfp4[tep4_cutedsl] accuracy/test_llm_api_pytorch.py::TestQwen3_5_397B_A17B::test_nvfp4[tep4_trtllm] accuracy/test_llm_api_pytorch.py::TestQwen3_5_4B::test_bf16 +accuracy/test_llm_api_pytorch.py::TestQwen3_5_4B::test_bf16_flashinfer accuracy/test_llm_api_pytorch.py::TestQwen3_5_4B::test_dflash accuracy/test_llm_api_pytorch.py::TestQwen3_5_4B::test_fp8 accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_dflash diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index 13f4bb62d60e..07114d7e870c 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -161,6 +161,7 @@ l0_h100: - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_dflash - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_dflash - accuracy/test_llm_api_pytorch.py::TestQwen3_5_4B::test_bf16 + - accuracy/test_llm_api_pytorch.py::TestQwen3_5_4B::test_bf16_flashinfer - accuracy/test_llm_api_pytorch.py::TestQwen3_5_4B::test_fp8 - accuracy/test_llm_api_pytorch.py::TestQwen3_5_4B::test_fp8_piecewise_cuda_graph - accuracy/test_llm_api_pytorch.py::TestQwen3_5_4B::test_dflash diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py index 3b519541fbf1..f60df6529a95 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py @@ -3,8 +3,9 @@ """Regression tests for Python, Cpp, and V2 Mamba cache managers.""" import os +from collections.abc import Callable from types import SimpleNamespace -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import pytest import torch @@ -12,6 +13,8 @@ from tensorrt_llm._torch.disaggregation.resource.kv_extractor import build_page_table_from_manager from tensorrt_llm._torch.disaggregation.resource.page import AttentionLayerGroup, MambaLayerGroup from tensorrt_llm._torch.disaggregation.transceiver import KvCacheTransceiverV2 +from tensorrt_llm._torch.flashinfer_utils import IS_FLASHINFER_AVAILABLE +from tensorrt_llm._torch.metadata import KVCacheParams from tensorrt_llm._torch.modules.mamba.mamba2_metadata import Mamba2Metadata from tensorrt_llm._torch.pyexecutor._util import ( KvCacheCreator, @@ -80,7 +83,11 @@ ) from tensorrt_llm.runtime.kv_cache_manager_v2 import KVCacheManager as RuntimeKVCacheManager +if IS_FLASHINFER_AVAILABLE: + from tensorrt_llm._torch.attention_backend import FlashInferAttentionMetadata + skip_no_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +skip_no_flashinfer = pytest.mark.skipif(not IS_FLASHINFER_AVAILABLE, reason="requires FlashInfer") def test_advance_replay_state_uses_checkpoint_predicate_and_skips_dummies(): @@ -2065,6 +2072,139 @@ def _build_v2_hybrid_with_mamba_layer( ) +@pytest.mark.parametrize( + "manager_factory", + [_build_hybrid_with_mamba_layer, _build_v2_hybrid_with_mamba_layer], +) +@skip_no_cuda +@skip_no_flashinfer +def test_flashinfer_hybrid_page_tables_use_attention_layer( + manager_factory: Callable[..., CppMambaHybridCacheManager | MambaHybridCacheManagerV2], +) -> None: + manager = manager_factory(max_batch_size=1) + try: + request_ids = [123] + manager.add_dummy_requests(request_ids, token_nums=[8], is_gen=False) + metadata = FlashInferAttentionMetadata( + seq_lens=torch.tensor([8], dtype=torch.int32), + num_contexts=1, + kv_cache_params=KVCacheParams( + use_cache=True, + num_cached_tokens_per_seq=[0], + ), + max_num_requests=1, + max_num_tokens=8, + kv_cache_manager=manager, + request_ids=request_ids, + ) + + with patch.object( + manager, + "get_batch_cache_indices_flat", + wraps=manager.get_batch_cache_indices_flat, + ) as get_indices: + metadata.prepare() + + assert metadata._get_pp_attention_layer_ids() == [1] + assert metadata._default_kv_layer_idx == 1 + assert get_indices.call_count >= 1 + assert all(call.kwargs.get("layer_idx") == 1 for call in get_indices.call_args_list) + assert torch.all(metadata.paged_kv_indices >= 0) + + with ( + patch.object( + manager, + "get_batch_cache_indices_flat", + wraps=manager.get_batch_cache_indices_flat, + ) as get_indices, + patch.object( + manager, + "get_batch_cache_indices", + wraps=manager.get_batch_cache_indices, + ) as get_block_ids, + ): + metadata._prepare_full_draft_page_table() + + assert get_indices.call_count >= 1 + assert get_block_ids.call_count >= 1 + assert all(call.kwargs.get("layer_idx") == 1 for call in get_indices.call_args_list) + assert all(call.kwargs.get("layer_idx") == 1 for call in get_block_ids.call_args_list) + finally: + manager.shutdown() + + +@skip_no_cuda +@skip_no_flashinfer +def test_flashinfer_metadata_falls_back_to_primary_pool_sizing() -> None: + manager = _build_v2_hybrid_with_mamba_layer(max_batch_size=1) + try: + with patch.object(manager, "get_buffers", None): + metadata = FlashInferAttentionMetadata( + seq_lens=torch.tensor([1], dtype=torch.int32), + num_contexts=1, + kv_cache_params=KVCacheParams( + use_cache=True, + num_cached_tokens_per_seq=[0], + ), + max_num_requests=1, + max_num_tokens=1, + kv_cache_manager=manager, + request_ids=[123], + ) + + assert metadata._paged_kv_indices.numel() == manager.blocks_in_primary_pool + finally: + manager.shutdown() + + +@pytest.mark.parametrize( + ("manager_factory", "manager_kwargs"), + [ + ( + _build_hybrid_with_mamba_layer, + { + "mamba_layer_mask": [True], + "attention_layer_mask": [False], + }, + ), + (_build_v2_hybrid_with_mamba_layer, {"num_attention_layers": 0}), + ], +) +@skip_no_cuda +@skip_no_flashinfer +def test_flashinfer_metadata_skips_paged_kv_on_recurrent_only_rank( + manager_factory: Callable[..., CppMambaHybridCacheManager | MambaHybridCacheManagerV2], + manager_kwargs: dict[str, object], +) -> None: + manager = manager_factory(max_batch_size=1, **manager_kwargs) + try: + request_ids = [123] + # Isolate FlashInfer's no-attention path from recurrent-state setup. + metadata = FlashInferAttentionMetadata( + seq_lens=torch.tensor([8], dtype=torch.int32), + num_contexts=1, + kv_cache_params=KVCacheParams( + use_cache=True, + num_cached_tokens_per_seq=[0], + ), + max_num_requests=1, + max_num_tokens=8, + kv_cache_manager=manager, + request_ids=request_ids, + mamba_metadata=False, + ) + + metadata.prepare() + + assert metadata._default_kv_layer_idx is None + assert metadata.paged_kv_indices.numel() == 0 + assert metadata.num_blocks == [0] + assert torch.count_nonzero(metadata.paged_kv_indptr) == 0 + assert torch.count_nonzero(metadata.paged_kv_last_page_len) == 0 + finally: + manager.shutdown() + + def _make_wide_spec_config(max_draft_len=2, tokens_per_gen_step=5): """Spec config whose per-step token width is wider than draft depth.