[None][fix] Forward layer_idx for VSWA + narrow FLASHINFER+spec rejection

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)
 

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]
 

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:

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

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(