From 631dc3d886b2f2632e32d6f91b920b33d18a9b71 Mon Sep 17 00:00:00 2001 From: Mihai Chiorean Date: Sat, 27 Jun 2026 17:05:34 -0700 Subject: [PATCH 01/18] [None][fix] Forward layer_idx to KV cache for FlashInfer / Vanilla backends (VSWA) Both FlashInfer and Vanilla attention backends called kv_cache_manager.get_batch_cache_indices() without a layer_idx argument. For VSWA models (e.g. Qwen3-Next family) the manager holds multiple KV pools with independent page numbering, so the argless call raises: ValueError: layer_idx or window_size must be provided for VSWA FlashInfer fix: detect _vswa_layer_to_pool at prepare() time and pass the primary pool's representative layer_idx to the initial block-count fetch. Per-pool indices are already rebuilt in the existing VSWA block below. Vanilla fix: prepare() detects kv_cache_manager.is_vswa and defers the block-ID lookup (setting block_ids_per_seq=None). forward() then calls get_batch_cache_indices(layer_idx=self.layer_idx) so each layer reads from the correct pool's page table. Signed-off-by: Mihai Chiorean --- .../_torch/attention_backend/flashinfer.py | 12 ++++++++- .../_torch/attention_backend/vanilla.py | 25 ++++++++++++++++--- 2 files changed, 32 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py index e6618779ed9c..70ce5b963cc0 100644 --- a/tensorrt_llm/_torch/attention_backend/flashinfer.py +++ b/tensorrt_llm/_torch/attention_backend/flashinfer.py @@ -890,8 +890,18 @@ def prepare(self) -> None: # indices of used cache blocks for each sequence assert self.request_ids is not None + # For VSWA models the KV cache manager has multiple pools with + # independent page numbering; calling get_batch_cache_indices without + # a layer_idx would raise ValueError. Use the primary pool's + # representative layer so the initial block-count computation has a + # valid window_size. Per-pool indices are rebuilt in the VSWA block + # below (lines ~944+) so this initial call is only used for num_blocks. + _vswa_init_layer: Optional[int] = None + if self._vswa_layer_to_pool is not None: + _primary_pool = self._vswa_layer_to_pool.get(0, 0) + _vswa_init_layer = self._vswa_pool_to_rep_layer.get(_primary_pool, 0) block_ids_per_seq = self.kv_cache_manager.get_batch_cache_indices( - self.request_ids) + self.request_ids, layer_idx=_vswa_init_layer) # number of tokens in the kv cache for each sequence in the batch cached_token_lens = torch.tensor( diff --git a/tensorrt_llm/_torch/attention_backend/vanilla.py b/tensorrt_llm/_torch/attention_backend/vanilla.py index a65f900dd95b..856418ab181b 100644 --- a/tensorrt_llm/_torch/attention_backend/vanilla.py +++ b/tensorrt_llm/_torch/attention_backend/vanilla.py @@ -62,10 +62,20 @@ class VanillaAttentionMetadata(AttentionMetadata): def prepare(self) -> None: super().prepare() - # indices of used cache blocks for each sequence + # indices of used cache blocks for each sequence. + # For VSWA models the correct pool depends on layer_idx, which is + # unknown at prepare() time (metadata is shared across all layers). + # Defer the lookup to VanillaAttention.forward() in that case. + if self.kv_cache_manager is None: + self.block_ids_per_seq = None + return assert self.request_ids is not None - self.block_ids_per_seq = self.kv_cache_manager.get_batch_cache_indices( - self.request_ids) if self.kv_cache_manager is not None else None + if getattr(self.kv_cache_manager, 'is_vswa', False): + # Signal to forward() that it must fetch per-layer block IDs. + self.block_ids_per_seq = None + else: + self.block_ids_per_seq = self.kv_cache_manager.get_batch_cache_indices( + self.request_ids) class VanillaAttention(AttentionBackend[VanillaAttentionMetadata]): @@ -503,8 +513,15 @@ def forward(self, attention_mask=forward_args.attention_mask) past_seen_tokens = metadata.kv_cache_params.num_cached_tokens_per_seq + # For VSWA models block_ids_per_seq is None (set by prepare()); fetch + # the correct pool's block IDs using this layer's index. + block_ids_per_seq = metadata.block_ids_per_seq + if block_ids_per_seq is None: + assert metadata.request_ids is not None + block_ids_per_seq = metadata.kv_cache_manager.get_batch_cache_indices( + metadata.request_ids, layer_idx=self.layer_idx) cache_indices = [ - block_ids[0] for block_ids in metadata.block_ids_per_seq + block_ids[0] for block_ids in block_ids_per_seq ] kv_cache_tensor = metadata.kv_cache_manager.get_buffers(self.layer_idx) From b504384185b5306672b68e411c94a34e6781a681 Mon Sep 17 00:00:00 2001 From: Mihai Chiorean Date: Sat, 27 Jun 2026 17:10:21 -0700 Subject: [PATCH 02/18] [None][fix] Lift FLASHINFER + speculative decoding rejection (vLLM-style) The blanket rejection at py_executor_creator.py:409 blocked ALL FlashInfer + one-engine speculative decoding (MTP) combinations, including the non-overlap-scheduler path that is actually viable. Root cause: BatchDecodeWithPagedKVCacheWrapper assumes q_len=1 per request. With MTP, generation requests carry num_nextn+1 tokens each, breaking the decode kernel. Fix (vLLM reorder_batch_threshold pattern adapted): - Add FlashInferAttentionMetadata._plan_mtp_gen_prefill(): plans a fresh BatchPrefillWithPagedKVCacheWrapper for the generation sub-batch using the rebased generation qo_indptr and paged_kv_indptr_decode slices. - In forward_impl(): detect MTP (gen seq_lens.max() > 1) and when not is_cuda_graph, route generation tokens through the prefill wrapper (mtp_gen_forward) instead of decode_forward. - Narrow the py_executor_creator guard to only reject when CUDA graphs are also enabled (where variable q_len truly is incompatible). CUDA-graph mode with FlashInfer + MTP remains unsupported (variable q_len per request cannot be captured). Non-graph mode (disable_overlap_ scheduler=True, no cuda_graph_config) now works end-to-end. Signed-off-by: Mihai Chiorean --- .../_torch/attention_backend/flashinfer.py | 96 ++++++++++++++++++- .../_torch/pyexecutor/py_executor_creator.py | 20 ++-- 2 files changed, 108 insertions(+), 8 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py index 70ce5b963cc0..9933802a88b9 100644 --- a/tensorrt_llm/_torch/attention_backend/flashinfer.py +++ b/tensorrt_llm/_torch/attention_backend/flashinfer.py @@ -407,6 +407,58 @@ def _do_plan_mla_decode(self, plan_params: MLAPlanParams) -> None: sm_scale=plan_params.sm_scale, ) + def _plan_mtp_gen_prefill(self, plan_params: PlanParams, + o_dtype: Optional[torch.dtype]) -> None: + """Plan a prefill-style wrapper for MTP generation tokens. + + Called when generation requests carry more than 1 token per sequence + (i.e. speculative decoding / MTP is active). + BatchDecodeWithPagedKVCacheWrapper assumes q_len=1 per sequence; for + MTP we route the generation sub-batch through a paged-prefill wrapper + instead. This path is NOT compatible with CUDA-graph capture (gated + on not self.is_cuda_graph at the call site). + """ + num_gen = self.num_generations + num_ctx = self.num_contexts + + # Rebase generation qo_indptr to start from 0. + gen_qo_indptr = ( + self._qo_indptr[num_ctx:num_ctx + num_gen + 1] - + self._qo_indptr[num_ctx]) + + gen_paged_kv_indptr = self.paged_kv_indptr_decode[:num_gen + 1] + gen_paged_kv_indices = self._paged_kv_indices[ + self.num_context_blocks:self.num_context_blocks + + self.num_generation_blocks] + gen_paged_kv_last_page = self._paged_kv_last_page_len[ + num_ctx:num_ctx + num_gen] + + # Allocate or reuse the MTP gen prefill wrapper and its indptr buffers. + # We allocate fresh each time rather than caching because batch shape + # can change run-to-run (and CUDA graph is disabled on this path). + self._mtp_gen_prefill_wrapper = flashinfer.BatchPrefillWithPagedKVCacheWrapper( + self.workspace_buffer, + self.kv_layout, + backend="fa2", + use_cuda_graph=False, + ) + self._mtp_gen_prefill_wrapper.plan( + gen_qo_indptr, + gen_paged_kv_indptr, + gen_paged_kv_indices, + gen_paged_kv_last_page, + plan_params.num_heads, + plan_params.num_kv_heads, + plan_params.head_dim, + self.page_size, + causal=True, + sm_scale=plan_params.sm_scale, + window_left=plan_params.window_left, + q_data_type=plan_params.q_dtype, + kv_data_type=plan_params.kv_dtype, + o_data_type=o_dtype, + ) + @property def paged_kv_indices(self) -> torch.Tensor: return self._paged_kv_indices[:self.num_generation_blocks + @@ -612,6 +664,17 @@ def _post_init_with_buffers(self, buffers) -> None: cache_name="_mla_kv_len_arr_buf", capture_graph=capture_graph, ) + # Buffers for MTP (speculative decoding) generation path. When + # generation requests carry more than 1 token per sequence the standard + # BatchDecodeWithPagedKVCacheWrapper cannot be used (it assumes q_len=1 + # per request). We instead route those through a separate prefill + # wrapper. Buffers are allocated here (not inside forward()) so their + # addresses are stable across calls; CUDA-graph capture is NOT supported + # on this path (gated on not is_cuda_graph in forward). + self._mtp_gen_qo_indptr_buf: Optional[torch.Tensor] = None + self._mtp_gen_paged_kv_indptr_buf: Optional[torch.Tensor] = None + self._mtp_gen_prefill_wrapper: Optional[ + flashinfer.BatchPrefillWithPagedKVCacheWrapper] = None # Rebind the wrapper to the freshly allocated buffers. self._ragged_prefill_wrapper = None self._mla_decode_wrapper = None @@ -1893,13 +1956,42 @@ def decode_forward(plan_params: PlanParams, out: torch.Tensor): attention_mask_data=effective_mask_data, flashinfer_backend=self.flashinfer_backend) + # Detect MTP: generation requests with >1 token per sequence. + # BatchDecodeWithPagedKVCacheWrapper assumes q_len=1; when MTP is + # active the generation sub-batch must use the prefill wrapper + # instead. Not compatible with CUDA-graph capture (variable + # q_len per request), so this path is gated on non-graph mode. + gen_seq_lens = metadata.seq_lens_cuda[ + num_contexts:num_contexts + num_generations] + is_mtp_gen = (num_generations > 0 + and not metadata.is_cuda_graph + and gen_seq_lens.max().item() > 1) + + def mtp_gen_forward(out: torch.Tensor): + """Run generation sub-batch through the paged prefill wrapper.""" + o_dtype = (torch.bfloat16 if plan_params.q_dtype in ( + torch.float8_e4m3fn, torch.float8_e5m2) else None) + metadata._plan_mtp_gen_prefill(plan_params, o_dtype) + assert metadata._mtp_gen_prefill_wrapper is not None + metadata._mtp_gen_prefill_wrapper.run( + q[num_ctx_tokens:].view(-1, self.num_heads, self.head_dim), + kv_cache, + out=out.view(-1, self.num_heads, self.head_dim), + ) + if num_contexts == 0: - decode_forward(plan_params, output) + if is_mtp_gen: + mtp_gen_forward(output) + else: + decode_forward(plan_params, output) elif num_generations == 0: prefill_forward(plan_params, output) else: prefill_forward(plan_params, output[:num_ctx_tokens, :]) - decode_forward(plan_params, output[num_ctx_tokens:, :]) + if is_mtp_gen: + mtp_gen_forward(output[num_ctx_tokens:, :]) + else: + decode_forward(plan_params, output[num_ctx_tokens:, :]) def forward(self, q: torch.Tensor, diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 797f2fd48666..0a464f942e26 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -406,14 +406,22 @@ def create_py_executor( ) llm_args.disable_overlap_scheduler = True - # Check FLASHINFER compatibility with one-engine speculative decoding - if llm_args.attn_backend == "FLASHINFER": + # FLASHINFER + overlap scheduler + one-engine spec-dec: not supported. + # The overlap scheduler pipelines decode and KV-append across steps; + # with multi-token generation requests (MTP/EAGLE) this conflicts with + # how FlashInfer plans its decode wrapper. Disable-overlap-scheduler + # users can use FLASHINFER + MTP without the overlap scheduler — the + # generation sub-batch is routed through the paged-prefill wrapper + # (see FlashInferAttentionMetadata._plan_mtp_gen_prefill). + # CUDA-graph mode is also unsupported for variable-q_len generation. + if (llm_args.attn_backend == "FLASHINFER" + and getattr(llm_args, "cuda_graph_config", None) is not None): raise ValueError( f"FLASHINFER attention backend is not supported with one-engine speculative " - f"decoding mode '{spec_config.spec_dec_mode.name}'. The FLASHINFER backend's " - f"decode path expects exactly 1 token per sequence, but one-engine speculative " - f"decoding requires multiple tokens per sequence. Please use 'TRTLLM' attention " - f"backend instead by setting attn_backend='TRTLLM'.") + f"decoding mode '{spec_config.spec_dec_mode.name}' when CUDA graphs are " + f"enabled (cuda_graph_config is set). The FlashInfer MTP generation path " + f"requires variable q_len per request which is incompatible with CUDA-graph " + f"capture. Either disable CUDA graphs or use 'TRTLLM' attention backend.") if mm_encoder_only: llm_args.mm_encoder_only = True From 464e91beb211fd6acdb416f6ab13f53e9b1dfa0a Mon Sep 17 00:00:00 2001 From: Mihai Chiorean Date: Sat, 27 Jun 2026 18:01:55 -0700 Subject: [PATCH 03/18] [None][fix] Handle V1 KVCacheManager VSWA in FlashInfer prepare() (follow-up) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The initial VSWA fix only handled KVCacheManagerV2 (checked via layer_to_pool_mapping_dict). KVCacheManagerV1 also exposes is_vswa=True when it has multiple distinct window sizes, but lacks the per-pool dict so _vswa_layer_to_pool stays None. The argless get_batch_cache_indices() call then raises ValueError: layer_idx or window_size must be provided. Add a fallback for V1: when is_vswa is True but _vswa_layer_to_pool is None, pick the first key from kv_cache_manager.layer_offsets and pass it as layer_idx. This resolves the window_size via the existing layer_offsets → max_attention_window_vec lookup path (V1 path in get_batch_cache_indices lines 1338-1342). Signed-off-by: Mihai Chiorean --- tensorrt_llm/_torch/attention_backend/flashinfer.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py index 9933802a88b9..1cff46306ae6 100644 --- a/tensorrt_llm/_torch/attention_backend/flashinfer.py +++ b/tensorrt_llm/_torch/attention_backend/flashinfer.py @@ -959,10 +959,22 @@ def prepare(self) -> None: # representative layer so the initial block-count computation has a # valid window_size. Per-pool indices are rebuilt in the VSWA block # below (lines ~944+) so this initial call is only used for num_blocks. + # For VSWA (V2) _vswa_layer_to_pool maps any layer to its pool; for + # VSWA (V1) that dict is None but kv_cache_manager.is_vswa is still + # True. In both cases we need any valid layer_idx to resolve the + # window_size. Use the first key in layer_offsets (guaranteed + # non-empty for any model with attention layers). _vswa_init_layer: Optional[int] = None if self._vswa_layer_to_pool is not None: + # V2 VSWA: use primary pool representative layer. _primary_pool = self._vswa_layer_to_pool.get(0, 0) _vswa_init_layer = self._vswa_pool_to_rep_layer.get(_primary_pool, 0) + elif getattr(self.kv_cache_manager, 'is_vswa', False): + # V1 VSWA: no per-pool mapping; any layer index resolves + # window_size via layer_offsets → max_attention_window_vec. + _layer_offsets = getattr(self.kv_cache_manager, 'layer_offsets', {}) + if _layer_offsets: + _vswa_init_layer = next(iter(_layer_offsets)) block_ids_per_seq = self.kv_cache_manager.get_batch_cache_indices( self.request_ids, layer_idx=_vswa_init_layer) From 5ccb9f311e016462e708c2053142a60067ec24cc Mon Sep 17 00:00:00 2001 From: Mihai Chiorean Date: Sat, 27 Jun 2026 18:08:13 -0700 Subject: [PATCH 04/18] [None][fix] Fix VSWA detection in FlashInfer prepare() for mixed-window models The V1 VSWA detection checked kv_cache_manager.is_vswa, but is_vswa has an additional 'all(w > 0)' guard that excludes models with a full-attention sentinel window (value -2147483647). Such models have len(max_attention_window_vec) > 1 but is_vswa = False, so the V1 fallback never fired. Fix: use len(max_attention_window_vec) > 1 directly as the trigger, matching the exact condition in get_batch_cache_indices that raises. This covers all V1 multi-window configurations regardless of sentinel values. Signed-off-by: Mihai Chiorean --- tensorrt_llm/_torch/attention_backend/flashinfer.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py index 1cff46306ae6..e239b2502bab 100644 --- a/tensorrt_llm/_torch/attention_backend/flashinfer.py +++ b/tensorrt_llm/_torch/attention_backend/flashinfer.py @@ -969,9 +969,10 @@ def prepare(self) -> None: # V2 VSWA: use primary pool representative layer. _primary_pool = self._vswa_layer_to_pool.get(0, 0) _vswa_init_layer = self._vswa_pool_to_rep_layer.get(_primary_pool, 0) - elif getattr(self.kv_cache_manager, 'is_vswa', False): - # V1 VSWA: no per-pool mapping; any layer index resolves - # window_size via layer_offsets → max_attention_window_vec. + elif len(getattr(self.kv_cache_manager, 'max_attention_window_vec', [])) > 1: + # V1 manager (or any manager without per-pool dict) with multiple + # window sizes: get_batch_cache_indices requires layer_idx to + # resolve the window_size. Use any valid layer index. _layer_offsets = getattr(self.kv_cache_manager, 'layer_offsets', {}) if _layer_offsets: _vswa_init_layer = next(iter(_layer_offsets)) From a01b7d4a204cd139cedc66927fb124ca007c181c Mon Sep 17 00:00:00 2001 From: Mihai Chiorean Date: Sat, 27 Jun 2026 23:29:22 -0700 Subject: [PATCH 05/18] [None][fix] Set blocks_in_primary_pool in VSWA / linear-attention branch (closes FIXME) The VSWA / linear-attention branch of KVCacheManager.__init__ computed blocks_per_window (window_size -> (primary, secondary)) but never set self.blocks_in_primary_pool or self.blocks_in_secondary_pool, leaving them unset for any non-estimation run on: - Pure-VSWA models (e.g. Qwen2.5-1M-style sliding-window-only) - Hybrid Mamba2 models using CppMambaHybridCacheManager Downstream, FlashInferAttentionMetadata.__post_init__ reads kv_cache_manager.blocks_in_primary_pool, raising AttributeError. Fix: after blocks_per_window is finalised (including the optional MPI allreduce), set blocks_in_primary_pool / blocks_in_secondary_pool to the maximum across all windows. FlashInfer already uses layer_idx for per-pool dispatch; the scalar attribute only needs to represent the upper bound (largest window), which is the correct semantic. Empirically validated on Spark (trtllm-smoke, Qwen3.6-35B-A3B NVFP4, attn_backend=FLASHINFER): the KV cache manager now initialises cleanly with two pools reported: primary blocks=8 [window=-2147483647], primary blocks=257 [window=4096] The fix resolves the FIXME added at this line. Signed-off-by: Mihai Chiorean --- tensorrt_llm/_torch/pyexecutor/resource_manager.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index ebb8fe753991..dc5bb0ee2393 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -406,9 +406,10 @@ 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. + # flashinfer.py accesses kv_cache_manager.blocks_in_primary_pool. + # For VSWA / linear-attention paths this is now set below to the + # maximum primary-pool size across all windows (largest window), + # which is what FlashInfer needs; per-call dispatch uses layer_idx. 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 @@ -494,6 +495,13 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int], logger.info( f"[MPI rank={mapping.rank}] Reduced blocks_per_window: {blocks_per_window}" ) + # Expose the largest-window primary/secondary block count as the + # scalar attribute consumed by FlashInfer (and any backend that + # calls kv_cache_manager.blocks_in_primary_pool directly). + self.blocks_in_primary_pool = max( + p for p, _ in blocks_per_window.values()) + self.blocks_in_secondary_pool = max( + s for _, s in blocks_per_window.values()) else: # Standard case: use original Python implementation self.blocks_in_primary_pool, self.blocks_in_secondary_pool = self.calculate_max_num_blocks( From cc0aa43f0755be380ac5ec1a6819d52e7d7ae75f Mon Sep 17 00:00:00 2001 From: Mihai Chiorean Date: Sun, 28 Jun 2026 09:32:11 -0700 Subject: [PATCH 06/18] [None][fix] Populate Mamba2Metadata.query_start_loc on pure-decode path The pure-decode branch (num_contexts == 0) of Mamba2Metadata.prepare() left query_start_loc as None, on the assumption that downstream callers only used query_start_loc_long. GatedDeltaNetMixer violates that: it always packs mamba_metadata.query_start_loc into kwargs and slices it, which raises TypeError when None. Reuse the existing int32 _arange_buffer (sized max_batch_size + 1) for the all-decode case. Numerically identical to what cu_seqlens[:batch_size+1] would be when every sequence contributes one token ([0, 1, ..., batch_size]). Surfaces on FlashInfer + Qwen3-Next (hybrid Mamba2) where the pure-decode / CUDA-graph capture path is first exercised. Non-FlashInfer paths reach the gdn_mixer with mixed batches where num_contexts > 0, so query_start_loc was always populated. Signed-off-by: Mihai Chiorean --- tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py index 5cab7283f033..08dff39d1b88 100644 --- a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py +++ b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py @@ -434,7 +434,7 @@ def prepare(self, attn_metadata: AttentionMetadata): self.chunk_indices = None self.chunk_offsets = None else: - self.query_start_loc = None + self.query_start_loc = self._arange_buffer[:batch_size + 1] self.query_start_loc_long = self._arange_buffer_long[:batch_size + 1] From 8d9087728e3d7ace6653034e57bc4ec58e0b5b9f Mon Sep 17 00:00:00 2001 From: Mihai Chiorean Date: Sun, 28 Jun 2026 10:00:03 -0700 Subject: [PATCH 07/18] [None][fix] Add spec_decoding_* fields to FlashInferAttentionMetadata Spec-decoding code (eagle3._prepare_attn_metadata_for_spec_dec, drafting_loops) reads and writes spec_decoding_packed_mask and friends on attn_metadata. These fields were declared only on TrtllmAttentionMetadata, not on FlashInferAttentionMetadata or the AttentionMetadata base. This was invisible while FLASHINFER + spec_dec was blanket-rejected. Now that the gate has been narrowed to one-engine + CUDA-graph only, FlashInfer + MTP exercises the code path and AttributeErrors. Schema-only fix: add the missing Optional[torch.Tensor] = None fields to FlashInferAttentionMetadata so the code path runs. Actually plumbing the packed mask into BatchPrefillWithPagedKVCacheWrapper.plan(custom_mask=...) for correct numerics is a separate concern and not addressed here. Signed-off-by: Mihai Chiorean --- tensorrt_llm/_torch/attention_backend/flashinfer.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py index e239b2502bab..f456e7b02344 100644 --- a/tensorrt_llm/_torch/attention_backend/flashinfer.py +++ b/tensorrt_llm/_torch/attention_backend/flashinfer.py @@ -169,6 +169,19 @@ class FlashInferAttentionMetadata(AttentionMetadata): _multi_item_params: Optional[FlashInferMultiItemParams] = field( init=False, default=None) + # Speculative-decoding metadata. Declared on TrtllmAttentionMetadata + # (trtllm.py) and read/written by the spec-dec code paths + # (eagle3._prepare_attn_metadata_for_spec_dec / drafting_loops). Now that + # the FLASHINFER + spec_dec gate has been narrowed to one-engine + CUDA + # graph only, FlashInfer + MTP exercises this code path and would + # AttributeError without these fields. Schema-only parity here; actually + # plumbing the packed mask into BatchPrefillWithPagedKVCacheWrapper.plan + # (custom_mask=...) for correct numerics is a separate concern. + spec_decoding_position_offsets: Optional[torch.Tensor] = None + spec_decoding_position_offsets_cpp: Optional[torch.Tensor] = None + spec_decoding_packed_mask: Optional[torch.Tensor] = None + spec_decoding_generation_lengths: Optional[torch.Tensor] = None + def needs_plan(self, plan_params: PlanParams) -> bool: if plan_params not in self._plan_params_to_wrappers: return True From 3bc069b37a0a2b939c43ce9370459b635b1d790c Mon Sep 17 00:00:00 2001 From: Mihai Chiorean Date: Sun, 28 Jun 2026 17:22:01 -0700 Subject: [PATCH 08/18] [None][fix] Disable separate draft KV cache for non-TRTLLM attn backends For one-engine speculative decoding (MTP, Eagle3) on hybrid Mamba models such as Qwen3-Next / Qwen3.6-35B-A3B-NVFP4, the separate-draft-KV-cache layout places the MTP draft layer (layer_idx == num_hidden_layers) in its own KVCacheManager. During the draft forward pass, the worker swaps that draft manager onto attn_metadata via SpecConfig.draft_kv_cache_context (and prepare_attn_metadata_for_draft_replay). Both helpers, however, short-circuit on `isinstance(attn_metadata, TrtllmAttentionMetadata)`, so when attn_backend != "TRTLLM" the swap never happens. The draft attention call then resolves metadata.kv_cache_manager.get_buffers(layer_idx=num_hidden_layers) against the target manager, whose layer_offsets only covers the dense range 0..(num_hidden_layers-1), and crashes with `KeyError`. Symptom on Spark: Qwen3.6-35B-A3B-NVFP4 + attn_backend=FLASHINFER + MTPDecodingConfig(num_nextn_predict_layers=1) raises KeyError: 40 at resource_manager.py:1435 (`layer_offset = self.layer_offsets[layer_idx]`) inside flashinfer.py:1853. Fix: for non-TRTLLM backends, refuse to create the separate draft KV cache manager. `extract_mamba_kv_cache_params` already extends the hybrid layer masks with MTP/draft attention entries when `layer_mask` is None, and `get_pp_layers` extends pp_layers symmetrically, so the combined target+draft layout places the draft layer in the target manager's layer_offsets and get_buffers resolves correctly. TRTLLM- attn retains the separate layout it already supports. Signed-off-by: Mihai Chiorean --- tensorrt_llm/_torch/pyexecutor/_util.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index ac0feab0347d..ca05a2dc1a0a 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -881,6 +881,28 @@ def _should_create_separate_draft_kv_cache(self) -> bool: "Attention DP is enabled, separate draft KV cache is not supported." ) return False + # The separate draft KV cache manager is swapped onto attn_metadata + # during the draft forward via SpecConfig.draft_kv_cache_context (and + # prepare_attn_metadata_for_draft_replay). Both helpers early-return + # for any attn_metadata that is not TrtllmAttentionMetadata, so non- + # TRTLLM attention backends (FlashInfer, FlashAttention, Vanilla) + # would keep using the target KV cache manager during the draft step + # — whose layer_offsets do not contain the draft layer index — and + # crash with KeyError in get_buffers(). Fall back to the combined + # cache layout (draft layers live in the target manager via the + # spec-layer extension in get_pp_layers / extract_mamba_kv_cache_params) + # for those backends. TRTLLM-attn keeps the separate layout it + # already supports. + if self._model_engine is not None: + attn_backend = getattr(self._model_engine.model.model_config, + 'attn_backend', None) + if attn_backend is not None and attn_backend != "TRTLLM": + logger.info( + f"Separate draft KV cache is not supported with " + f"attn_backend={attn_backend} (only TRTLLM supports the " + f"draft-KV-cache swap on attn_metadata); falling back to " + f"combined target+draft KV cache layout.") + return False return should_use_separate_draft_kv_cache(self._speculative_config) def _get_effective_draft_config(self) -> ModelConfig: From 8f4970747f954c0d51d850bad718e42599de1594 Mon Sep 17 00:00:00 2001 From: Mihai Chiorean Date: Sun, 28 Jun 2026 20:54:49 -0700 Subject: [PATCH 09/18] [None][fix] Pick non-linear pool for FlashInfer _vswa_init_layer On Qwen3-Next-family hybrid Mamba models, layer 0 is a Mamba2 / linear- attention layer. FlashInferAttentionMetadata.prepare() was resolving _vswa_init_layer to the rep layer of pool 0 - which is the linear / recurrent pool. BlockManager::getFreeBlock returns NEGATIVE placeholder block IDs for Mamba recurrent-state slots. Those negatives populated _paged_kv_indices and leaked into _vswa_pool_buf_0 via the primary_buf copy. On the second autotuner warmup (post-estimation pool resize), swap_paged_kv_indices_for_layer for a full-attn layer surfaced stale linear-pool data, append_paged_kv_cache received negative kv_indices and crashed with an illegal memory access. Fix: in _vswa_init_layer resolution, prefer pools whose window > 0 (real SWA / full-attention) over linear pools (window == -INT_MAX). The fallback for the all-linear edge case is preserved. Same heuristic applied to the V1 KVCacheManager fallback added in 464e91beb2. Signed-off-by: Mihai Chiorean --- .../_torch/attention_backend/flashinfer.py | 72 +++++++++++++++++-- 1 file changed, 67 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py index f456e7b02344..1a1d94fbe29e 100644 --- a/tensorrt_llm/_torch/attention_backend/flashinfer.py +++ b/tensorrt_llm/_torch/attention_backend/flashinfer.py @@ -477,6 +477,56 @@ def paged_kv_indices(self) -> torch.Tensor: return self._paged_kv_indices[:self.num_generation_blocks + self.num_context_blocks] + def _pool_window_for_layer(self, layer_idx: int) -> Optional[int]: + """Return the per-pool window size for *layer_idx*, or None if unknown. + + Linear / recurrent (Mamba2) layers use a negative INT_MAX sentinel + for window; real SWA / full-attention layers have window > 0 + (or None, which V2 represents as max_seq_len). Callers use this + to avoid picking a linear-pool layer as the primary pool, since + those pools return negative placeholder block IDs that would + crash append_paged_kv_cache with an illegal memory access. + """ + mgr = self.kv_cache_manager + vec = getattr(mgr, 'max_attention_window_vec', None) + if not vec: + return None + layer_offsets = getattr(mgr, 'layer_offsets', None) + if layer_offsets is None or layer_idx not in layer_offsets: + return None + off = layer_offsets[layer_idx] + # V2 manager: has layer_to_pool_mapping_dict; vec is indexed + # per-layer (length matches num_layers when user supplies + # per-layer max_attention_window, e.g. Qwen3-Next hybrid). + if hasattr(mgr, 'layer_to_pool_mapping_dict') and 0 <= off < len(vec): + w = vec[off] + return w if w is not None else (1 << 31) + # V1 manager: vec is per-pool; resolve pool via _layer_to_pool_idx. + l2p = getattr(mgr, '_layer_to_pool_idx', None) + if l2p is not None: + pool_idx = l2p.get(off) + if pool_idx is not None and 0 <= pool_idx < len(vec): + w = vec[pool_idx] + return w if w is not None else (1 << 31) + return None + + def _pick_vswa_primary_pool_id(self) -> int: + """Pick a non-linear pool as the FlashInfer 'primary' pool. + + On hybrid Mamba models (e.g. Qwen3-Next) layer 0 is a linear / + recurrent layer; its pool returns NEGATIVE placeholder block IDs + from BlockManager::getFreeBlock that crash append_paged_kv_cache. + Pick the first pool whose representative layer has window > 0. + Fall back to the layer-0 pool if no full-attention pool exists. + """ + assert self._vswa_layer_to_pool is not None + default_pool = self._vswa_layer_to_pool.get(0, 0) + for pool_id, rep_layer in self._vswa_pool_to_rep_layer.items(): + w = self._pool_window_for_layer(rep_layer) + if w is not None and w > 0: + return pool_id + return default_pool + def get_paged_kv_indices_for_layer(self, layer_idx: int) -> torch.Tensor: """Return page indices for the pool that *layer_idx* belongs to. @@ -979,16 +1029,26 @@ def prepare(self) -> None: # non-empty for any model with attention layers). _vswa_init_layer: Optional[int] = None if self._vswa_layer_to_pool is not None: - # V2 VSWA: use primary pool representative layer. - _primary_pool = self._vswa_layer_to_pool.get(0, 0) + # V2 VSWA: prefer the rep layer of a non-linear (window > 0) + # pool. On hybrid Mamba models layer 0 belongs to the linear + # pool, whose block IDs are negative placeholders that would + # poison _paged_kv_indices and crash append_paged_kv_cache. + _primary_pool = self._pick_vswa_primary_pool_id() _vswa_init_layer = self._vswa_pool_to_rep_layer.get(_primary_pool, 0) elif len(getattr(self.kv_cache_manager, 'max_attention_window_vec', [])) > 1: # V1 manager (or any manager without per-pool dict) with multiple # window sizes: get_batch_cache_indices requires layer_idx to - # resolve the window_size. Use any valid layer index. + # resolve the window_size. Prefer a layer whose pool window is + # positive (skip linear / recurrent pools whose window sentinel + # is negative); fall back to the first layer if none qualifies. _layer_offsets = getattr(self.kv_cache_manager, 'layer_offsets', {}) if _layer_offsets: _vswa_init_layer = next(iter(_layer_offsets)) + for _lid in _layer_offsets: + _w = self._pool_window_for_layer(_lid) + if _w is not None and _w > 0: + _vswa_init_layer = _lid + break block_ids_per_seq = self.kv_cache_manager.get_batch_cache_indices( self.request_ids, layer_idx=_vswa_init_layer) @@ -1042,7 +1102,9 @@ def prepare(self) -> None: # 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 must be non-linear so _paged_kv_indices / the + # primary buffer hold positive page IDs (see _pick_vswa_primary_pool_id). + primary_pool_id = self._pick_vswa_primary_pool_id() # 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. @@ -1178,7 +1240,7 @@ def prepare(self) -> None: # 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._pick_vswa_primary_pool_id() 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) From 34ca84c1bfb066291e1ba9b4cd5609db5cd5c927 Mon Sep 17 00:00:00 2001 From: Mihai Chiorean Date: Mon, 29 Jun 2026 11:16:08 -0700 Subject: [PATCH 10/18] [None][fix] Detect hybrid Mamba pools in FlashInfer metadata Signed-off-by: Mihai Chiorean --- tensorrt_llm/_torch/attention_backend/flashinfer.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py index 1a1d94fbe29e..9dc0da9ffc15 100644 --- a/tensorrt_llm/_torch/attention_backend/flashinfer.py +++ b/tensorrt_llm/_torch/attention_backend/flashinfer.py @@ -670,10 +670,14 @@ def _post_init_with_buffers(self, buffers) -> None: capture_graph=capture_graph, ) - # Detect VSWA: check if the manager has multiple pools. - # Guard on layer_to_pool_mapping_dict which is V2-specific — V1 - # managers also expose is_vswa but lack the per-pool infrastructure. - if (getattr(self.kv_cache_manager, 'is_vswa', False) and hasattr( + # Detect multi-pool managers. kv_cache_manager.is_vswa excludes + # hybrid Mamba managers because their recurrent-state pool uses a + # negative sentinel window, but FlashInfer still needs the V2 pool + # mapping to avoid feeding those placeholder block IDs to paged KV. + has_multiple_windows = len( + getattr(self.kv_cache_manager, 'max_attention_window_vec', + [])) > 1 + if (has_multiple_windows and hasattr( self.kv_cache_manager, 'layer_to_pool_mapping_dict')): mgr = self.kv_cache_manager self._vswa_layer_to_pool = {} From c498619cad37e76f0b06423131d31a2fd4cc0c58 Mon Sep 17 00:00:00 2001 From: Mihai Chiorean Date: Mon, 29 Jun 2026 11:23:31 -0700 Subject: [PATCH 11/18] [None][fix] Build FlashInfer pool map from layer offsets Signed-off-by: Mihai Chiorean --- .../_torch/attention_backend/flashinfer.py | 86 +++++++++++-------- 1 file changed, 48 insertions(+), 38 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py index 9dc0da9ffc15..e7fbc50762d6 100644 --- a/tensorrt_llm/_torch/attention_backend/flashinfer.py +++ b/tensorrt_llm/_torch/attention_backend/flashinfer.py @@ -677,45 +677,55 @@ def _post_init_with_buffers(self, buffers) -> None: has_multiple_windows = len( getattr(self.kv_cache_manager, 'max_attention_window_vec', [])) > 1 - if (has_multiple_windows and hasattr( - self.kv_cache_manager, 'layer_to_pool_mapping_dict')): + if has_multiple_windows: mgr = self.kv_cache_manager - self._vswa_layer_to_pool = {} - self._vswa_pool_to_rep_layer: Dict[int, int] = {} - for layer_idx in getattr(mgr, 'layer_offsets', {}): - layer_offset = mgr.layer_offsets[layer_idx] - pool_id = mgr.layer_to_pool_mapping_dict[layer_offset] - self._vswa_layer_to_pool[layer_idx] = pool_id - if pool_id not in self._vswa_pool_to_rep_layer: - self._vswa_pool_to_rep_layer[pool_id] = layer_idx - # Build head_dim → pool_id mapping using V2 per-layer head_dim - self._vswa_head_dim_to_pool: Dict[int, int] = {} - if hasattr(mgr, 'head_dim_per_layer'): - for layer_idx, pool_id in self._vswa_layer_to_pool.items(): - hd = mgr.head_dim_per_layer[ - mgr.layer_offsets[layer_idx]] - if hd not in self._vswa_head_dim_to_pool: - self._vswa_head_dim_to_pool[hd] = pool_id - - # Pre-allocate VSWA pool cache buffers. These must be - # stable (never reallocated) so that CUDA-graph-recorded - # copies reference valid addresses across replays. - # Use the maximum page count across ALL pools (not just the - # primary) so that secondary pool buffers are large enough. - all_pool_pages = max_num_pages - 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: - all_pool_pages = max(all_pool_pages, lbuf.shape[0]) - for pool_id in set(self._vswa_layer_to_pool.values()): - buf_key = f'_vswa_pool_buf_{pool_id}' - if getattr(self, buf_key, None) is None: - setattr( - self, buf_key, - torch.empty(all_pool_pages, - dtype=torch.int, - device='cuda')) + layer_to_pool_mapping = getattr(mgr, + 'layer_to_pool_mapping_dict', + None) + layer_to_pool_idx = getattr(mgr, '_layer_to_pool_idx', None) + if layer_to_pool_mapping is not None or layer_to_pool_idx is not None: + self._vswa_layer_to_pool = {} + self._vswa_pool_to_rep_layer: Dict[int, int] = {} + for layer_idx in getattr(mgr, 'layer_offsets', {}): + layer_offset = mgr.layer_offsets[layer_idx] + if layer_to_pool_mapping is not None: + pool_id = layer_to_pool_mapping[layer_offset] + else: + pool_id = layer_to_pool_idx.get(layer_offset) + if pool_id is None: + continue + self._vswa_layer_to_pool[layer_idx] = pool_id + if pool_id not in self._vswa_pool_to_rep_layer: + self._vswa_pool_to_rep_layer[pool_id] = layer_idx + # Build head_dim → pool_id mapping using V2 per-layer head_dim + self._vswa_head_dim_to_pool: Dict[int, int] = {} + if hasattr(mgr, 'head_dim_per_layer'): + for layer_idx, pool_id in self._vswa_layer_to_pool.items(): + hd = mgr.head_dim_per_layer[ + mgr.layer_offsets[layer_idx]] + if hd not in self._vswa_head_dim_to_pool: + self._vswa_head_dim_to_pool[hd] = pool_id + + # Pre-allocate VSWA pool cache buffers. These must be + # stable (never reallocated) so that CUDA-graph-recorded + # copies reference valid addresses across replays. + # Use the maximum page count across ALL pools (not just the + # primary) so that secondary pool buffers are large enough. + all_pool_pages = max_num_pages + 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: + all_pool_pages = max(all_pool_pages, + lbuf.shape[0]) + for pool_id in set(self._vswa_layer_to_pool.values()): + buf_key = f'_vswa_pool_buf_{pool_id}' + if getattr(self, buf_key, None) is None: + setattr( + self, buf_key, + torch.empty(all_pool_pages, + dtype=torch.int, + device='cuda')) # Stable buffers for FlashInfer MLA decode; required for CUDA graphs. self._mla_qo_indptr_buf = self.get_empty( buffers, From f640dc67b9e17e4301cce657e62c944fffad5dd3 Mon Sep 17 00:00:00 2001 From: Mihai Chiorean Date: Mon, 29 Jun 2026 11:26:36 -0700 Subject: [PATCH 12/18] [None][fix] Skip recurrent pools when sizing FlashInfer buffers Signed-off-by: Mihai Chiorean --- tensorrt_llm/_torch/attention_backend/flashinfer.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py index e7fbc50762d6..f2aa3b35a4b0 100644 --- a/tensorrt_llm/_torch/attention_backend/flashinfer.py +++ b/tensorrt_llm/_torch/attention_backend/flashinfer.py @@ -714,6 +714,9 @@ def _post_init_with_buffers(self, buffers) -> None: all_pool_pages = max_num_pages if hasattr(self.kv_cache_manager, 'layer_offsets'): for lid in self.kv_cache_manager.layer_offsets: + window = self._pool_window_for_layer(lid) + if window is not None and window <= 0: + continue lbuf = self.kv_cache_manager.get_buffers(lid) if lbuf is not None: all_pool_pages = max(all_pool_pages, From 2a446207b28ffb1d1f0814c6cf209af8e3506fbf Mon Sep 17 00:00:00 2001 From: Mihai Chiorean Date: Mon, 29 Jun 2026 11:29:47 -0700 Subject: [PATCH 13/18] [None][fix] Skip zero-KV layers when sizing FlashInfer buffers Signed-off-by: Mihai Chiorean --- tensorrt_llm/_torch/attention_backend/flashinfer.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py index f2aa3b35a4b0..daf601d88dd0 100644 --- a/tensorrt_llm/_torch/attention_backend/flashinfer.py +++ b/tensorrt_llm/_torch/attention_backend/flashinfer.py @@ -714,8 +714,9 @@ def _post_init_with_buffers(self, buffers) -> None: all_pool_pages = max_num_pages if hasattr(self.kv_cache_manager, 'layer_offsets'): for lid in self.kv_cache_manager.layer_offsets: - window = self._pool_window_for_layer(lid) - if window is not None and window <= 0: + layer_offset = self.kv_cache_manager.layer_offsets[lid] + if (self.kv_cache_manager. + num_kv_heads_per_layer[layer_offset] == 0): continue lbuf = self.kv_cache_manager.get_buffers(lid) if lbuf is not None: From cbf46d55ce895c2d1d2d591f8a905a2c801ab6ee Mon Sep 17 00:00:00 2001 From: Mihai Chiorean Date: Mon, 29 Jun 2026 11:34:10 -0700 Subject: [PATCH 14/18] [None][fix] Allocate FlashInfer buffers for all KV pools Signed-off-by: Mihai Chiorean --- .../_torch/attention_backend/flashinfer.py | 37 +++++++++++++++++-- 1 file changed, 34 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py index daf601d88dd0..d48bdd9b0d3a 100644 --- a/tensorrt_llm/_torch/attention_backend/flashinfer.py +++ b/tensorrt_llm/_torch/attention_backend/flashinfer.py @@ -488,10 +488,19 @@ def _pool_window_for_layer(self, layer_idx: int) -> Optional[int]: crash append_paged_kv_cache with an illegal memory access. """ mgr = self.kv_cache_manager + layer_offsets = getattr(mgr, 'layer_offsets', None) + if layer_offsets is not None and layer_idx in layer_offsets: + layer_offset = layer_offsets[layer_idx] + l2p = getattr(mgr, '_layer_to_pool_idx', None) + if l2p is not None: + pool_idx = l2p.get(layer_offset) + if pool_idx is not None: + window = self._pool_window_for_pool_id(pool_idx) + if window is not None: + return window vec = getattr(mgr, 'max_attention_window_vec', None) if not vec: return None - layer_offsets = getattr(mgr, 'layer_offsets', None) if layer_offsets is None or layer_idx not in layer_offsets: return None off = layer_offsets[layer_idx] @@ -510,6 +519,20 @@ def _pool_window_for_layer(self, layer_idx: int) -> Optional[int]: return w if w is not None else (1 << 31) return None + def _pool_window_for_pool_id(self, pool_id: int) -> Optional[int]: + """Return the configured window for a KV manager pool.""" + mgr = self.kv_cache_manager + pool_configurations = getattr(mgr, 'pool_configurations', None) + if pool_configurations is not None and 0 <= pool_id < len( + pool_configurations): + window = pool_configurations[pool_id].window_size + return window if window is not None else (1 << 31) + vec = getattr(mgr, 'max_attention_window_vec', None) + if vec is not None and 0 <= pool_id < len(vec): + window = vec[pool_id] + return window if window is not None else (1 << 31) + return None + def _pick_vswa_primary_pool_id(self) -> int: """Pick a non-linear pool as the FlashInfer 'primary' pool. @@ -522,7 +545,9 @@ def _pick_vswa_primary_pool_id(self) -> int: assert self._vswa_layer_to_pool is not None default_pool = self._vswa_layer_to_pool.get(0, 0) for pool_id, rep_layer in self._vswa_pool_to_rep_layer.items(): - w = self._pool_window_for_layer(rep_layer) + w = self._pool_window_for_pool_id(pool_id) + if w is None: + w = self._pool_window_for_layer(rep_layer) if w is not None and w > 0: return pool_id return default_pool @@ -722,7 +747,13 @@ def _post_init_with_buffers(self, buffers) -> None: if lbuf is not None: all_pool_pages = max(all_pool_pages, lbuf.shape[0]) - for pool_id in set(self._vswa_layer_to_pool.values()): + pool_ids = set(self._vswa_layer_to_pool.values()) + num_pools = getattr(self.kv_cache_manager, 'num_pools', + None) + if num_pools is not None: + pool_ids.update(range(num_pools)) + pool_ids.add(self._pick_vswa_primary_pool_id()) + for pool_id in pool_ids: buf_key = f'_vswa_pool_buf_{pool_id}' if getattr(self, buf_key, None) is None: setattr( From 852ec182b47686d807f4fd3f023e0280e352dccd Mon Sep 17 00:00:00 2001 From: Mihai Chiorean Date: Mon, 29 Jun 2026 11:46:34 -0700 Subject: [PATCH 15/18] [None][fix] Map hybrid attention layers to KV pool Signed-off-by: Mihai Chiorean --- .../_torch/attention_backend/flashinfer.py | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py index d48bdd9b0d3a..f4210db9d0ed 100644 --- a/tensorrt_llm/_torch/attention_backend/flashinfer.py +++ b/tensorrt_llm/_torch/attention_backend/flashinfer.py @@ -711,6 +711,14 @@ def _post_init_with_buffers(self, buffers) -> None: if layer_to_pool_mapping is not None or layer_to_pool_idx is not None: self._vswa_layer_to_pool = {} self._vswa_pool_to_rep_layer: Dict[int, int] = {} + positive_pool_ids = [ + pool_id for pool_id in range( + getattr(mgr, 'num_pools', 0)) + if (self._pool_window_for_pool_id(pool_id) is not None + and self._pool_window_for_pool_id(pool_id) > 0) + ] + single_positive_pool_id = (positive_pool_ids[0] if len( + positive_pool_ids) == 1 else None) for layer_idx in getattr(mgr, 'layer_offsets', {}): layer_offset = mgr.layer_offsets[layer_idx] if layer_to_pool_mapping is not None: @@ -718,7 +726,13 @@ def _post_init_with_buffers(self, buffers) -> None: else: pool_id = layer_to_pool_idx.get(layer_offset) if pool_id is None: - continue + if single_positive_pool_id is None: + continue + pool_id = single_positive_pool_id + pool_window = self._pool_window_for_pool_id(pool_id) + if ((pool_window is None or pool_window <= 0) + and single_positive_pool_id is not None): + pool_id = single_positive_pool_id self._vswa_layer_to_pool[layer_idx] = pool_id if pool_id not in self._vswa_pool_to_rep_layer: self._vswa_pool_to_rep_layer[pool_id] = layer_idx From 98e5ee8449b0c742d567bd6cd7c85119e0620ef0 Mon Sep 17 00:00:00 2001 From: Mihai Chiorean Date: Mon, 29 Jun 2026 11:53:21 -0700 Subject: [PATCH 16/18] [None][fix] Fallback to KV pool for FlashInfer hybrid layers Signed-off-by: Mihai Chiorean --- .../_torch/attention_backend/flashinfer.py | 35 ++++++++++++++----- 1 file changed, 27 insertions(+), 8 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py index f4210db9d0ed..f2d756b3e6c2 100644 --- a/tensorrt_llm/_torch/attention_backend/flashinfer.py +++ b/tensorrt_llm/_torch/attention_backend/flashinfer.py @@ -533,6 +533,23 @@ def _pool_window_for_pool_id(self, pool_id: int) -> Optional[int]: return window if window is not None else (1 << 31) return None + def _positive_vswa_pool_ids(self) -> list[int]: + """Return KV cache pool ids that can back FlashInfer paged KV.""" + mgr = self.kv_cache_manager + num_pools = getattr(mgr, 'num_pools', None) + if num_pools is None: + pool_configurations = getattr(mgr, 'pool_configurations', None) + if pool_configurations is not None: + num_pools = len(pool_configurations) + else: + num_pools = len(getattr(mgr, 'max_attention_window_vec', [])) + positive_pool_ids: list[int] = [] + for pool_id in range(num_pools): + window = self._pool_window_for_pool_id(pool_id) + if window is not None and window > 0: + positive_pool_ids.append(pool_id) + return positive_pool_ids + def _pick_vswa_primary_pool_id(self) -> int: """Pick a non-linear pool as the FlashInfer 'primary' pool. @@ -544,6 +561,9 @@ def _pick_vswa_primary_pool_id(self) -> int: """ assert self._vswa_layer_to_pool is not None default_pool = self._vswa_layer_to_pool.get(0, 0) + positive_pool_ids = self._positive_vswa_pool_ids() + if positive_pool_ids: + return positive_pool_ids[0] for pool_id, rep_layer in self._vswa_pool_to_rep_layer.items(): w = self._pool_window_for_pool_id(pool_id) if w is None: @@ -580,7 +600,7 @@ def swap_paged_kv_indices_for_layer(self, layer_idx: int) -> None: return pool_id = self._vswa_layer_to_pool.get(layer_idx) if pool_id is None: - return # Layer not in VSWA mapping + pool_id = self._pick_vswa_primary_pool_id() active = getattr(self, '_vswa_active_pool_id', None) if pool_id == active and not self.is_cuda_graph: return # Buffer already has the right data @@ -711,12 +731,7 @@ def _post_init_with_buffers(self, buffers) -> None: if layer_to_pool_mapping is not None or layer_to_pool_idx is not None: self._vswa_layer_to_pool = {} self._vswa_pool_to_rep_layer: Dict[int, int] = {} - positive_pool_ids = [ - pool_id for pool_id in range( - getattr(mgr, 'num_pools', 0)) - if (self._pool_window_for_pool_id(pool_id) is not None - and self._pool_window_for_pool_id(pool_id) > 0) - ] + positive_pool_ids = self._positive_vswa_pool_ids() single_positive_pool_id = (positive_pool_ids[0] if len( positive_pool_ids) == 1 else None) for layer_idx in getattr(mgr, 'layer_offsets', {}): @@ -1097,7 +1112,11 @@ def prepare(self) -> None: # pool, whose block IDs are negative placeholders that would # poison _paged_kv_indices and crash append_paged_kv_cache. _primary_pool = self._pick_vswa_primary_pool_id() - _vswa_init_layer = self._vswa_pool_to_rep_layer.get(_primary_pool, 0) + _vswa_init_layer = self._vswa_pool_to_rep_layer.get(_primary_pool) + if _vswa_init_layer is None: + _layer_offsets = getattr(self.kv_cache_manager, 'layer_offsets', + {}) + _vswa_init_layer = next(iter(_layer_offsets), 0) elif len(getattr(self.kv_cache_manager, 'max_attention_window_vec', [])) > 1: # V1 manager (or any manager without per-pool dict) with multiple # window sizes: get_batch_cache_indices requires layer_idx to From 98daf477e5ec01493c5b50eee41ba2d40717796f Mon Sep 17 00:00:00 2001 From: Mihai Chiorean Date: Mon, 29 Jun 2026 11:59:41 -0700 Subject: [PATCH 17/18] [None][fix] Seed FlashInfer indices from attention layer Signed-off-by: Mihai Chiorean --- .../_torch/attention_backend/flashinfer.py | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py index f2d756b3e6c2..1e47e61ec069 100644 --- a/tensorrt_llm/_torch/attention_backend/flashinfer.py +++ b/tensorrt_llm/_torch/attention_backend/flashinfer.py @@ -550,6 +550,19 @@ def _positive_vswa_pool_ids(self) -> list[int]: positive_pool_ids.append(pool_id) return positive_pool_ids + def _first_positive_window_layer_id(self) -> int: + """Return a layer id that uses a real attention KV window.""" + mgr = self.kv_cache_manager + layer_offsets = getattr(mgr, 'layer_offsets', {}) + vec = getattr(mgr, 'max_attention_window_vec', None) + if vec: + pattern_len = len(vec) + for layer_idx in layer_offsets: + window = vec[layer_idx % pattern_len] + if window is not None and window > 0: + return layer_idx + return next(iter(layer_offsets), 0) + def _pick_vswa_primary_pool_id(self) -> int: """Pick a non-linear pool as the FlashInfer 'primary' pool. @@ -1114,9 +1127,7 @@ def prepare(self) -> None: _primary_pool = self._pick_vswa_primary_pool_id() _vswa_init_layer = self._vswa_pool_to_rep_layer.get(_primary_pool) if _vswa_init_layer is None: - _layer_offsets = getattr(self.kv_cache_manager, 'layer_offsets', - {}) - _vswa_init_layer = next(iter(_layer_offsets), 0) + _vswa_init_layer = self._first_positive_window_layer_id() elif len(getattr(self.kv_cache_manager, 'max_attention_window_vec', [])) > 1: # V1 manager (or any manager without per-pool dict) with multiple # window sizes: get_batch_cache_indices requires layer_idx to From 4af66682724a83bcc07252b05904fba22895e68f Mon Sep 17 00:00:00 2001 From: Mihai Chiorean Date: Mon, 29 Jun 2026 12:06:31 -0700 Subject: [PATCH 18/18] [None][fix] Add FlashInfer request type metadata Signed-off-by: Mihai Chiorean --- tensorrt_llm/_torch/attention_backend/flashinfer.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py index 1e47e61ec069..22732d7def77 100644 --- a/tensorrt_llm/_torch/attention_backend/flashinfer.py +++ b/tensorrt_llm/_torch/attention_backend/flashinfer.py @@ -143,6 +143,8 @@ class FlashInferAttentionMetadata(AttentionMetadata): _cached_token_lens: torch.Tensor = field(init=False) _plan_params_to_wrappers: Dict[PlanParams, FlashInferWrappers] = field(init=False) + host_request_types: torch.Tensor = field(init=False) + host_request_types_runtime: torch.Tensor = field(init=False) # MLA wrappers and stable buffers. # Cached plan params + is-planned flag let prepare() refresh the plan @@ -703,6 +705,10 @@ def _post_init_with_buffers(self, buffers) -> None: self._cached_token_lens = torch.empty((self.max_num_requests, ), dtype=torch.int, device='cuda') + self.host_request_types = torch.empty((self.max_num_requests, ), + dtype=torch.int, + pin_memory=prefer_pinned(), + device='cpu') self._batch_indices = torch.empty((self.max_num_tokens, ), dtype=torch.int, device='cuda') @@ -1078,6 +1084,10 @@ def prepare(self) -> None: dim=0, dtype=torch.int32, out=self._qo_indptr[1:self.seq_lens_cuda.size(0) + 1]) + num_seqs = self.num_contexts + self.num_generations + self.host_request_types[:self.num_contexts].fill_(0) + self.host_request_types[self.num_contexts:num_seqs].fill_(1) + self.host_request_types_runtime = self.host_request_types[:num_seqs] if self.multi_item_part_lens is not None: self._multi_item_params = self._process_multi_item_part_lens(