[None][fix] Select FlashInfer hybrid default KV layer

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.

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

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"

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

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

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