From f32e1976448c44ecec4dc7230822d137bb93d628 Mon Sep 17 00:00:00 2001 From: Mihai Chiorean Date: Thu, 6 Aug 2026 15:41:15 -0700 Subject: [PATCH 1/4] [None][fix] Select FlashInfer hybrid primary KV layer Signed-off-by: Mihai Chiorean --- .../_torch/attention_backend/flashinfer.py | 83 +++++++++++----- .../executor/test_mamba_cache_manager.py | 98 ++++++++++++++++++- 2 files changed, 156 insertions(+), 25 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py index d49ab084a128..31da5772b56c 100644 --- a/tensorrt_llm/_torch/attention_backend/flashinfer.py +++ b/tensorrt_llm/_torch/attention_backend/flashinfer.py @@ -702,7 +702,8 @@ 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) + primary_pool_id = self._vswa_layer_to_pool.get( + target._primary_kv_layer_idx, 0) self._vswa_active_pool_id = primary_pool_id self._host_paged_kv_indices = self._host_pool_indices[ primary_pool_id] @@ -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._primary_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._primary_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_primary_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_attention_layer_ids(self) -> list[int]: + """Return local layers backed by an attention KV cache.""" + 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_primary_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._primary_kv_layer_idx is not None + return self.kv_cache_manager.get_batch_cache_indices_flat( + request_ids, + num_blocks, + layer_idx=self._primary_kv_layer_idx, + ) + def _post_init_with_buffers(self, buffers) -> None: capture_graph = self.is_cuda_graph + self._primary_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,22 @@ 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 + attention_layer_ids = self._get_attention_layer_ids() + self._primary_kv_layer_idx = next(iter(attention_layer_ids), None) + max_num_blocks = int( + getattr(self.kv_cache_manager, "blocks_in_primary_pool", 0)) + for layer_idx in attention_layer_ids: + layer_buffer = self.kv_cache_manager.get_buffers(layer_idx) + if layer_buffer is not None: + max_num_blocks = max(max_num_blocks, layer_buffer.shape[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 +976,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 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 +1426,15 @@ def _to_int32_tensor(arr: np.ndarray) -> torch.Tensor: assert self.request_ids is not None + if self._primary_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._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 +1453,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_primary_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,14 +1471,14 @@ 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) + primary_pool_id = self._vswa_layer_to_pool.get( + self._primary_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) + primary_buf[:total_idx].copy_(paged_kv_indices, non_blocking=True) self._vswa_pool_indices_cache = { primary_pool_id: primary_buf, } @@ -1574,7 +1608,8 @@ def _to_int32_tensor(arr: np.ndarray) -> torch.Tensor: # VSWA: restore primary pool indices as the default. 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) + primary_pool_id = self._vswa_layer_to_pool.get( + self._primary_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] self._paged_kv_indices[:total_blocks].copy_(src, non_blocking=True) diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py index 8ceddf7e03a3..f9355a3f6727 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(): @@ -1741,6 +1748,95 @@ 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._primary_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_skips_paged_kv_on_recurrent_only_rank() -> None: + manager = _build_v2_hybrid_with_mamba_layer(max_batch_size=1, num_attention_layers=0) + 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, + ) + + metadata.prepare() + + assert metadata._primary_kv_layer_idx is None + assert metadata.paged_kv_indices.numel() == 0 + assert metadata.num_blocks == [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. From 907755952bb2dd43c90cc4d3837e1b8049332686 Mon Sep 17 00:00:00 2001 From: Mihai Chiorean Date: Tue, 11 Aug 2026 23:56:56 -0700 Subject: [PATCH 2/4] [None][fix] Address FlashInfer hybrid KV review Signed-off-by: Mihai Chiorean --- .../_torch/attention_backend/flashinfer.py | 72 +++++++++---------- .../_torch/pyexecutor/resource_manager.py | 3 - .../defs/accuracy/test_llm_api_pytorch.py | 14 ++++ .../executor/test_mamba_cache_manager.py | 5 +- 4 files changed, 53 insertions(+), 41 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py index 31da5772b56c..1a16b1e5889f 100644 --- a/tensorrt_llm/_torch/attention_backend/flashinfer.py +++ b/tensorrt_llm/_torch/attention_backend/flashinfer.py @@ -702,13 +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( - target._primary_kv_layer_idx, 0) - self._vswa_active_pool_id = primary_pool_id + default_pool_id = self._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 @@ -771,15 +771,15 @@ 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 - or self._primary_kv_layer_idx is None): + 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, layer_idx=self._primary_kv_layer_idx) + 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._get_primary_kv_indices(self.request_ids, + 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) @@ -838,8 +838,8 @@ def __post_init__(self) -> None: super().__post_init__() self._post_init_with_buffers(self.cuda_graph_buffers) - def _get_attention_layer_ids(self) -> list[int]: - """Return local layers backed by an attention KV cache.""" + 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", @@ -851,21 +851,21 @@ def _get_attention_layer_ids(self) -> list[int]: if num_kv_heads[layer_offset] > 0 ] - def _get_primary_kv_indices(self, request_ids: list[int], + 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._primary_kv_layer_idx 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._primary_kv_layer_idx, + layer_idx=self._default_kv_layer_idx, ) def _post_init_with_buffers(self, buffers) -> None: capture_graph = self.is_cuda_graph - self._primary_kv_layer_idx: Optional[int] = None + 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 @@ -947,11 +947,11 @@ 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: - attention_layer_ids = self._get_attention_layer_ids() - self._primary_kv_layer_idx = next(iter(attention_layer_ids), None) - max_num_blocks = int( - getattr(self.kv_cache_manager, "blocks_in_primary_pool", 0)) - for layer_idx in attention_layer_ids: + pp_attention_layer_ids = self._get_pp_attention_layer_ids() + self._default_kv_layer_idx = next(iter(pp_attention_layer_ids), + None) + max_num_blocks = 0 + for layer_idx in pp_attention_layer_ids: layer_buffer = self.kv_cache_manager.get_buffers(layer_idx) if layer_buffer is not None: max_num_blocks = max(max_num_blocks, layer_buffer.shape[0]) @@ -976,7 +976,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 attention_layer_ids: + 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)) @@ -1426,7 +1426,7 @@ def _to_int32_tensor(arr: np.ndarray) -> torch.Tensor: assert self.request_ids is not None - if self._primary_kv_layer_idx is 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 @@ -1453,7 +1453,7 @@ 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._get_primary_kv_indices(self.request_ids, + paged_kv_indices = self._get_default_kv_indices(self.request_ids, self.num_blocks) self._paged_kv_indices[:paged_kv_indices.size(0)].copy_( @@ -1471,20 +1471,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( - self._primary_kv_layer_idx, 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_(paged_kv_indices, 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 = \ @@ -1495,7 +1495,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. @@ -1605,17 +1605,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( - self._primary_kv_layer_idx, 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 eb3d4193bd8c..a174b2ac317d 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -449,9 +449,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 09c0224f26af..c22faae28fc5 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -6218,6 +6218,20 @@ def test_bf16(self): task.evaluate(llm, extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS) + @skip_pre_hopper + def test_bf16_flashinfer(self): + 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/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py index f9355a3f6727..b4c9cbf55c5f 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py @@ -1781,7 +1781,8 @@ def test_flashinfer_hybrid_page_tables_use_attention_layer( ) as get_indices: metadata.prepare() - assert metadata._primary_kv_layer_idx == 1 + 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) @@ -1830,7 +1831,7 @@ def test_flashinfer_metadata_skips_paged_kv_on_recurrent_only_rank() -> None: metadata.prepare() - assert metadata._primary_kv_layer_idx is None + assert metadata._default_kv_layer_idx is None assert metadata.paged_kv_indices.numel() == 0 assert metadata.num_blocks == [0] finally: From b30fdc12d1b404426a73933c4ab2baa4ac7362b3 Mon Sep 17 00:00:00 2001 From: Mihai Chiorean Date: Wed, 12 Aug 2026 00:14:02 -0700 Subject: [PATCH 3/4] [None][test] Register FlashInfer hybrid accuracy coverage Signed-off-by: Mihai Chiorean --- tests/integration/defs/accuracy/test_llm_api_pytorch.py | 2 +- tests/integration/test_lists/qa/llm_function_core.txt | 1 + tests/integration/test_lists/qa/llm_function_rtx6k.txt | 1 + tests/integration/test_lists/test-db/l0_h100.yml | 1 + 4 files changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index c22faae28fc5..745be72d7872 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -6219,7 +6219,7 @@ def test_bf16(self): extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS) @skip_pre_hopper - def test_bf16_flashinfer(self): + def test_bf16_flashinfer(self) -> None: model_path = f"{llm_models_root()}/Qwen3.5-4B" with LLM(model_path, trust_remote_code=True, diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index 13e07dfdae25..158db0025a4f 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -744,6 +744,7 @@ accuracy/test_llm_api_pytorch.py::TestQwen3_4B::test_eagle3 accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16[tp1-CUTLASS] accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16[tp1-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_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/integration/test_lists/qa/llm_function_rtx6k.txt b/tests/integration/test_lists/qa/llm_function_rtx6k.txt index 4260a5328d77..768688cf21f6 100644 --- a/tests/integration/test_lists/qa/llm_function_rtx6k.txt +++ b/tests/integration/test_lists/qa/llm_function_rtx6k.txt @@ -180,6 +180,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 fb57eb7c1a8f..0f4fdacf47eb 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -159,6 +159,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 From 5ddb185f14705f8f645015acee87261d5de42488 Mon Sep 17 00:00:00 2001 From: Mihai Chiorean Date: Wed, 12 Aug 2026 12:00:18 -0700 Subject: [PATCH 4/4] [None][fix] Harden FlashInfer hybrid metadata Signed-off-by: Mihai Chiorean --- .../_torch/attention_backend/flashinfer.py | 27 ++++++++-- .../executor/test_mamba_cache_manager.py | 49 +++++++++++++++++-- 2 files changed, 68 insertions(+), 8 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py index 1a16b1e5889f..a9f84ac23e40 100644 --- a/tensorrt_llm/_torch/attention_backend/flashinfer.py +++ b/tensorrt_llm/_torch/attention_backend/flashinfer.py @@ -702,7 +702,7 @@ 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 - default_pool_id = self._vswa_layer_to_pool.get( + 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[ @@ -950,11 +950,23 @@ def _post_init_with_buffers(self, buffers) -> None: 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 - for layer_idx in pp_attention_layer_ids: - layer_buffer = self.kv_cache_manager.get_buffers(layer_idx) - if layer_buffer is not None: - max_num_blocks = max(max_num_blocks, layer_buffer.shape[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, (max_num_blocks, ), @@ -1432,6 +1444,11 @@ def _to_int32_tensor(arr: np.ndarray) -> torch.Tensor: 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 diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py index b4c9cbf55c5f..b984579e2088 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py @@ -1811,11 +1811,51 @@ def test_flashinfer_hybrid_page_tables_use_attention_layer( @skip_no_cuda @skip_no_flashinfer -def test_flashinfer_metadata_skips_paged_kv_on_recurrent_only_rank() -> None: - manager = _build_v2_hybrid_with_mamba_layer(max_batch_size=1, num_attention_layers=0) +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] - manager.add_dummy_requests(request_ids, token_nums=[8], is_gen=False) + # Isolate FlashInfer's no-attention path from recurrent-state setup. metadata = FlashInferAttentionMetadata( seq_lens=torch.tensor([8], dtype=torch.int32), num_contexts=1, @@ -1827,6 +1867,7 @@ def test_flashinfer_metadata_skips_paged_kv_on_recurrent_only_rank() -> None: max_num_tokens=8, kv_cache_manager=manager, request_ids=request_ids, + mamba_metadata=False, ) metadata.prepare() @@ -1834,6 +1875,8 @@ def test_flashinfer_metadata_skips_paged_kv_on_recurrent_only_rank() -> None: 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()