diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py index a1fb669369a0..c16d2ffaee36 100644 --- a/tensorrt_llm/_torch/attention_backend/flashinfer.py +++ b/tensorrt_llm/_torch/attention_backend/flashinfer.py @@ -37,6 +37,7 @@ from tensorrt_llm.functional import AttentionMaskType from tensorrt_llm.logger import logger from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.runtime.kv_cache_manager_v2._common import BAD_PAGE_INDEX from ..metadata import KVCacheParams from ..utils import get_global_attrs, get_model_extra_attrs, torch_multi_arange @@ -232,6 +233,8 @@ class MLAPlanParams: @dataclass(kw_only=True) class FlashInferWrappers: is_planned: bool + fa2_plan_num_blocks: Optional[tuple[int, ...]] = field(default=None, + repr=False) decode_wrapper: Optional[ flashinfer.BatchDecodeWithPagedKVCacheWrapper] = None prefill_wrapper: Optional[ @@ -579,6 +582,20 @@ def get_paged_kv_indices_for_layer(self, layer_idx: int) -> torch.Tensor: total_blocks = self.num_generation_blocks + self.num_context_blocks return self._vswa_pool_indices_cache[pool_id][:total_blocks] + def _sanitize_swa_page_indices(self, page_indices: torch.Tensor, + layer_idx: int) -> None: + """Replace evicted SWA pages with a safe in-range page index.""" + window_vec = getattr(self.kv_cache_manager, 'max_attention_window_vec', + None) + if not window_vec or window_vec[layer_idx % len(window_vec)] is None: + return + + # KVCacheManagerV2 marks evicted out-of-window pages with -1. + # FlashInfer may dereference page IDs before applying window_left, so + # keep masked positions in range. The SWA mask excludes their values + # from the attention result. + page_indices.masked_fill_(page_indices == BAD_PAGE_INDEX, 0) + def swap_paged_kv_indices_for_layer(self, layer_idx: int) -> None: """Copy pool-specific page indices into the shared buffer. @@ -1425,12 +1442,32 @@ def _clean_cached_plans(self, *, defer_plan: bool): # corresponding forward pass. So, flush them out here as they won't be relevant for # subsequent forward calls. if plan_params.attention_mask_data is None and plan_params.multi_item_params is None: - self._plan_params_to_wrappers[plan_params].is_planned = False + wrappers = self._plan_params_to_wrappers[plan_params] + if wrappers.fa2_plan_num_blocks is not None: + continue + wrappers.is_planned = False if not defer_plan: self._plan_with_params(plan_params) else: del self._plan_params_to_wrappers[plan_params] + def _refresh_fa2_cuda_graph_plans(self) -> None: + """Refresh captured FA2 schedules after page metadata is finalized.""" + num_blocks = tuple(self.num_blocks[self.num_contexts:]) + for plan_params, wrappers in self._plan_params_to_wrappers.items(): + if (plan_params.attention_mask_data is not None + or plan_params.multi_item_params is not None + or wrappers.fa2_plan_num_blocks is None): + continue + if not num_blocks: + wrappers.fa2_plan_num_blocks = None + elif wrappers.fa2_plan_num_blocks != num_blocks: + # Graph replay does not re-enter forward_impl. Each wrapper + # owns its persistent integer plan workspace, while the shared + # float workspace is run scratch. + wrappers.is_planned = False + self._plan_with_params(plan_params) + def prepare(self) -> None: def _to_int32_tensor(arr: np.ndarray) -> torch.Tensor: @@ -1542,8 +1579,18 @@ 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 + primary_layer_idx = None + if self._vswa_layer_to_pool is not None: + primary_pool_id = self._vswa_layer_to_pool.get(0, 0) + primary_layer_idx = self._vswa_pool_to_rep_layer[primary_pool_id] + else: + layer_offsets = getattr(self.kv_cache_manager, 'layer_offsets', {}) + primary_layer_idx = next(iter(layer_offsets), None) + paged_kv_indices = self.kv_cache_manager.get_batch_cache_indices_flat( - self.request_ids, self.num_blocks) + self.request_ids, self.num_blocks, layer_idx=primary_layer_idx) + if primary_layer_idx is not None: + self._sanitize_swa_page_indices(paged_kv_indices, primary_layer_idx) self._paged_kv_indices[:paged_kv_indices.size(0)].copy_( paged_kv_indices, non_blocking=True) @@ -1579,6 +1626,7 @@ def _to_int32_tensor(arr: np.ndarray) -> torch.Tensor: pool_indices = \ self.kv_cache_manager.get_batch_cache_indices_flat( self.request_ids, self.num_blocks, layer_idx=rep_layer) + self._sanitize_swa_page_indices(pool_indices, rep_layer) buf = getattr(self, f'_vswa_pool_buf_{pool_id}') buf[:pool_indices.size(0)].copy_(pool_indices, non_blocking=True) @@ -1656,19 +1704,6 @@ def _to_int32_tensor(arr: np.ndarray) -> torch.Tensor: self._positions[:positions.size(0)].copy_(positions, non_blocking=True) - # Multi-wrapper case (Gemma4 hybrid: different head_dim per layer) - # shares one workspace_buffer; eager plan() would overwrite earlier - # wrappers' workspace, so defer plan() to forward_impl. Single-wrapper - # case (e.g., Llama, Gemma3 uniform head_dim) needs eager plan() here - # because forward_impl cannot plan() during cuda-graph stream capture. - active_wrappers = [ - pp for pp in self._plan_params_to_wrappers - if pp.attention_mask_data is None - ] - defer_plan = len(active_wrappers) > 1 - if not (self._is_separate_kv_draft_view and self.is_cuda_graph): - self._clean_cached_plans(defer_plan=defer_plan) - # Re-plan MLA wrappers outside of forward/capture using the params # cached by prior warmup forwards. Forward still handles first-use or # dtype/shape changes by syncing only on a plan cache miss. @@ -1798,6 +1833,20 @@ def _to_int32_tensor(arr: np.ndarray) -> torch.Tensor: non_blocking=True) if self.num_generations < batch_size: kv_lens_buf[self.num_generations:batch_size].zero_() + + # Refresh captured FA2 schedules only after all page metadata updates. + # Defer ordinary multi-wrapper plans to forward_impl; single-wrapper + # models still plan eagerly because forward_impl cannot plan during + # graph capture. + active_wrappers = [ + pp for pp in self._plan_params_to_wrappers + if pp.attention_mask_data is None + ] + defer_plan = len(active_wrappers) > 1 + if not (self._is_separate_kv_draft_view and self.is_cuda_graph): + self._refresh_fa2_cuda_graph_plans() + self._clean_cached_plans(defer_plan=defer_plan) + if (not self._is_shared_kv_draft_view and not self._is_separate_kv_draft_view and self._draft_metadata is not None): @@ -1987,8 +2036,12 @@ def prefill_plan(): custom_mask=plan_params.attention_mask_data, ) + use_graph_tensor_cores = self.is_cuda_graph and plan_params.head_dim > 128 if wrappers.decode_wrapper is None: use_tensor_cores = self._use_tensor_cores(plan_params) + # Gemma4's H256/H512 plans need a tensor-core wrapper with a stable + # CUDA Graph launch layout. prepare() may refresh its split-K + # schedule in the wrapper's fixed workspace as KV pages change. wrappers.decode_wrapper = \ flashinfer.BatchDecodeWithPagedKVCacheWrapper( @@ -1998,11 +2051,11 @@ def prefill_plan(): paged_kv_indptr_buffer=self.paged_kv_indptr_decode, paged_kv_indices_buffer=self._paged_kv_indices, paged_kv_last_page_len_buffer=self._paged_kv_last_page_len, - use_tensor_cores=use_tensor_cores + use_tensor_cores=use_tensor_cores or use_graph_tensor_cores or flashinfer_backend == "trtllm-gen", backend=flashinfer_backend if flashinfer_backend != "fa2" else - ("fa2" if torch.cuda.get_device_capability(0) == ( + ("fa2" if torch.cuda.get_device_capability() == ( 9, 0) else "auto"), ) decode_wrapper = wrappers.decode_wrapper @@ -2042,7 +2095,11 @@ def decode_plan(): block_tables=block_tables, # Keep FlashInfer's recorded graph shape aligned with the wrapper cache key. q_len_per_req=plan_params.q_len_per_req, + disable_split_kv=False, ) + if use_graph_tensor_cores and decode_wrapper._backend == 'fa2': + wrappers.fa2_plan_num_blocks = tuple( + self.num_blocks[self.num_contexts:]) self._publish_decode_wrapper_kv_lens(decode_wrapper) # Must sync after append_paged_kv_cache and before plan. diff --git a/tensorrt_llm/_torch/models/modeling_gemma4.py b/tensorrt_llm/_torch/models/modeling_gemma4.py index 9acd5a5e2b3f..f535b17e53bf 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma4.py +++ b/tensorrt_llm/_torch/models/modeling_gemma4.py @@ -1339,11 +1339,10 @@ def get_context_mask( """Build context mask with causal + bidirectional for MM tokens. Returns a [extend_len, prefix_len + extend_len] mask where: - - The first `prefix_len` columns (cached/paged history) are True for - all rows. SWA window enforcement is delegated to the kernel's - window_left clip. Bidirectional MM across the prefix/extend - boundary is NOT supported here; callers must ensure chunk - boundaries do not split a multimodal block. + - The first `prefix_len` columns (cached/paged history) apply the + sliding window using absolute token positions. Bidirectional MM + across the prefix/extend boundary is NOT supported here; callers + must ensure chunk boundaries do not split a multimodal block. - The last `extend_len` columns follow the original causal + (optional) sliding window + MM-bidirectional logic. """ @@ -1360,9 +1359,19 @@ def get_context_mask( causal_mask = causal_mask.masked_fill(token_type_mask, True) if prefix_len > 0: - prefix_block = torch.ones( - extend_len, prefix_len, dtype=causal_mask.dtype, device=device - ) + if ( + effective_sliding_window is not None + and effective_sliding_window < prefix_len + extend_len + ): + query_pos = prefix_len + pos + prefix_pos = torch.arange(prefix_len, device=device) + prefix_block = ( + prefix_pos.unsqueeze(0) > query_pos.unsqueeze(1) - effective_sliding_window + ) + else: + prefix_block = torch.ones( + extend_len, prefix_len, dtype=causal_mask.dtype, device=device + ) causal_mask = torch.cat([prefix_block, causal_mask], dim=1) return causal_mask diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index 79d77502baf4..e62cbbc3f9bc 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -68,6 +68,7 @@ l0_h100: - unittest/_torch/modeling -k "modeling_mixtral" - unittest/_torch/modeling -k "modeling_gemma3" - unittest/_torch/modeling -k "modeling_gpt_oss" + - unittest/_torch/modeling -k "modeling_gemma4" - unittest/_torch/modeling -k "modeling_whisper" # CPU-only log-mel parity - unittest/_torch/modeling/test_modeling_nemotron_h.py::test_nemotron_h_sanity # Real-weight Nano CG/overlap and chunked-prefill path smoke (MoE L0 cannot diff --git a/tests/unittest/_torch/modeling/test_modeling_gemma4.py b/tests/unittest/_torch/modeling/test_modeling_gemma4.py index 9745126f0ee2..e57baacb94c2 100644 --- a/tests/unittest/_torch/modeling/test_modeling_gemma4.py +++ b/tests/unittest/_torch/modeling/test_modeling_gemma4.py @@ -48,6 +48,7 @@ from tensorrt_llm._utils import is_sm_100f from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig from tensorrt_llm.mapping import Mapping +from tensorrt_llm.runtime.kv_cache_manager_v2._common import BAD_PAGE_INDEX if TYPE_CHECKING: from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 @@ -896,6 +897,17 @@ def test_assistant_uses_target_kv_sources(self): "attention_k_eq_v": True, } +# 12B-real-dims: GQA=2 sliding (16/8), GQA=16 full K=V (16/1), +# hd=256/512. +GEMMA4_12B_REAL_DIMS_CONFIG = { + **GEMMA4_E2B_REAL_DIMS_CONFIG, + "num_hidden_layers": 12, + "num_attention_heads": 16, + "num_key_value_heads": 8, + "num_global_key_value_heads": 1, + "attention_k_eq_v": True, +} + # 26B-real-dims: GQA=2 sliding (16/8), GQA=2 full K=V (16/8), hd=256/512. GEMMA4_26B_REAL_DIMS_CONFIG = { **GEMMA4_E2B_REAL_DIMS_CONFIG, @@ -912,6 +924,7 @@ def _build_gemma4_kv_cache_manager( num_blocks=4, tokens_per_block=32, batch_size=1, + enable_swa_eviction: bool = False, ): """Create KVCacheManagerV2 supporting Gemma4 per-layer head_dim / kv_heads. @@ -962,17 +975,18 @@ def _build_gemma4_kv_cache_manager( # Set per-layer max_attention_window when head_dim or kv_heads differ # across layers, so V2 creates separate pool groups for different page - # sizes. ``max_seq_len - 1`` on sliding layers prevents V2 block - # eviction that would cause FlashInfer page index OOB when kv_lens - # exceeds sliding_window. + # sizes. sliding_window = getattr(config, "sliding_window", None) max_attn_window = None needs_vswa = isinstance(head_dim, list) and len(set(head_dim)) > 1 if not needs_vswa: needs_vswa = isinstance(num_kv_heads, list) and len(set(num_kv_heads)) > 1 if needs_vswa and sliding_window: + swa_window = ( + min(sliding_window, max_seq_len - 1) if enable_swa_eviction else max_seq_len - 1 + ) max_attn_window = [ - max_seq_len - 1 if lt == "sliding_attention" else max_seq_len for lt in layer_types + swa_window if lt == "sliding_attention" else max_seq_len for lt in layer_types ] kv_cache_config = KvCacheConfigV2( @@ -2064,21 +2078,8 @@ def test_vswa_pool_cache_not_aliased(self): @unittest.mock.patch( "tensorrt_llm.runtime.kv_cache_manager_v2._utils.assert_critical", lambda *a, **kw: None ) - def test_vswa_no_eviction_with_long_sequence(self): - """VSWA: sliding pool must not evict blocks when max_attention_window - uses max_seq_len - 1 (the fix for page index OOB). - - Root cause: when _util.py used the model's sliding_window (e.g. 512) - as max_attention_window for sliding layers, V2 would evict old blocks - when kv_lens exceeded the window. But FlashInfer's prepare() computes - num_blocks from the FULL kv_lens, so the page indices for evicted - blocks become stale → illegal memory access. - - The fix uses max_seq_len - 1 instead of sliding_window, preventing - eviction while keeping is_vswa=True. This test verifies that with - the fix, a sequence longer than sliding_window still has all its - blocks allocated (no eviction) and page indices are within bounds. - """ + def test_vswa_evicted_page_indices_are_sanitized(self) -> None: + """FlashInfer metadata replaces evicted SWA page markers.""" from tensorrt_llm._torch.attention_backend.utils import get_attention_backend from tensorrt_llm._torch.metadata import KVCacheParams @@ -2087,44 +2088,45 @@ def test_vswa_no_eviction_with_long_sequence(self): config_dict["sliding_window"] = 64 config = Gemma4TextConfig(**config_dict) - # num_blocks=4 → max_seq_len = 4*128 = 512, much larger than - # sliding_window=64. With the fix, max_attention_window for sliding - # layers = 511 (max_seq_len - 1), so V2 won't evict. - kv_cache_manager = self._get_kv_cache_manager(config, num_blocks=4) + kv_cache_manager = self._get_kv_cache_manager( + config, num_blocks=4, enable_swa_eviction=True + ) - # Allocate a request with tokens > sliding_window + # Allocate a generation request longer than the sliding window. request_ids = [1] - token_nums = [128] # 128 tokens >> sliding_window (64) - kv_cache_manager.add_dummy_requests(request_ids, token_nums) + cached_tokens = 126 + token_nums = [cached_tokens + 1] + kv_cache_manager.add_dummy_requests(request_ids, token_nums, is_gen=True) + + num_blocks = (token_nums[0] + kv_cache_manager.tokens_per_block - 1) // ( + kv_cache_manager.tokens_per_block + ) + raw_indices = kv_cache_manager.get_batch_cache_indices_flat( + request_ids, [num_blocks], layer_idx=0 + ) + self.assertIn(BAD_PAGE_INDEX, raw_indices.tolist()) metadata_cls = get_attention_backend("FLASHINFER").Metadata metadata = metadata_cls( - seq_lens=torch.tensor([128], dtype=torch.int), - num_contexts=1, + seq_lens=torch.ones(1, dtype=torch.int), + num_contexts=0, kv_cache_params=KVCacheParams( use_cache=True, - num_cached_tokens_per_seq=[0], + num_cached_tokens_per_seq=[cached_tokens], ), max_num_requests=1, max_num_tokens=8192, kv_cache_manager=kv_cache_manager, request_ids=request_ids, - prompt_lens=[128], ) with torch.inference_mode(): metadata.prepare() - # num_blocks should be based on full kv_lens (128 tokens), - # not clamped to sliding_window (64 tokens). - expected_blocks = ( - 128 + kv_cache_manager.tokens_per_block - 1 - ) // kv_cache_manager.tokens_per_block - self.assertEqual( - metadata.num_blocks[0], - expected_blocks, - f"num_blocks should be {expected_blocks} (from full kv_lens=128), " - f"not clamped to sliding_window={config_dict['sliding_window']}", + self.assertEqual(metadata.num_blocks[0], num_blocks) + self.assertNotIn( + BAD_PAGE_INDEX, + metadata.get_paged_kv_indices_for_layer(0).cpu().tolist(), ) # Page indices must be within bounds for EVERY layer @@ -2335,6 +2337,31 @@ def test_bidirectional_mask_gating(self): # But causal for text tokens self.assertFalse(mask_26b[0, 1].item(), "Text token 0 should NOT attend to 1") + @torch.no_grad() + def test_chunked_context_mask_applies_prefix_window(self) -> None: + """Chunked prefill preserves the full-sequence sliding-window mask.""" + config_dict = deepcopy(GEMMA4_E4B_LIKE_CONFIG) + config_dict["use_bidirectional_attention"] = "vision" + config = Gemma4TextConfig(**config_dict) + model_config = ModelConfig(pretrained_config=config, attn_backend="FLASHINFER") + model = Gemma4ForCausalLM(model_config).to(config.torch_dtype).to("cuda") + + window = 64 + chunk_start = 48 + token_type_ids = torch.zeros(96, dtype=torch.long, device="cuda") + token_type_ids[64:80] = 1 + + full_mask = model.get_context_mask(token_type_ids, effective_sliding_window=window) + chunk_mask = model.get_context_mask( + token_type_ids[chunk_start:], + effective_sliding_window=window, + prefix_len=chunk_start, + ) + + torch.testing.assert_close(chunk_mask, full_mask[chunk_start:]) + self.assertFalse(chunk_mask[-1, 0].item()) + self.assertTrue(chunk_mask[-1, 32].item()) + @torch.no_grad() def test_bidirectional_mask_only_applies_to_sliding_layers(self): """Full-attention layers retain the standard causal mask.""" @@ -2660,6 +2687,7 @@ def _expected_decode_block_table( source_offset += page_count return expected + @unittest.skipUnless(is_sm_100f(), "trtllm-gen attention requires SM100f") @torch.no_grad() @unittest.mock.patch( "tensorrt_llm.runtime.kv_cache_manager_v2._utils.assert_critical", lambda *a, **kw: None @@ -2701,6 +2729,7 @@ def test_shared_kv_draft_view(self) -> None: rtol=0, ) + @unittest.skipUnless(is_sm_100f(), "trtllm-gen attention requires SM100f") @torch.no_grad() @unittest.mock.patch( "tensorrt_llm.runtime.kv_cache_manager_v2._utils.assert_critical", lambda *a, **kw: None @@ -2738,6 +2767,7 @@ def test_cuda_graph_trtllm_gen_block_table_transitions(self) -> None: self.assertEqual(wrappers.decode_block_table_active_rows, len(new_page_counts)) self.assertEqual(wrappers.decode_block_table_active_width, max(new_page_counts)) + @unittest.skipUnless(is_sm_100f(), "trtllm-gen attention requires SM100f") @torch.no_grad() @unittest.mock.patch( "tensorrt_llm.runtime.kv_cache_manager_v2._utils.assert_critical", lambda *a, **kw: None @@ -2788,6 +2818,7 @@ def test_cuda_graph_trtllm_gen_host_table_growth_keeps_device_pointer(self) -> N rtol=0, ) + @unittest.skipUnless(is_sm_100f(), "trtllm-gen attention requires SM100f") @torch.no_grad() @unittest.mock.patch( "tensorrt_llm.runtime.kv_cache_manager_v2._utils.assert_critical", lambda *a, **kw: None @@ -3213,11 +3244,12 @@ def test_cuda_graph_multi_step_decode(self): kv_cache_manager.shutdown() + @unittest.skipUnless(is_sm_100f(), "trtllm-gen attention requires SM100f") @torch.no_grad() @unittest.mock.patch( "tensorrt_llm.runtime.kv_cache_manager_v2._utils.assert_critical", lambda *a, **kw: None ) - def test_cuda_graph_decode_high_gqa(self): + def test_cuda_graph_decode_high_gqa(self) -> None: """CUDA graph decode with GQA=8 and real head_dim (E2B-like). Uses E2B real-dims config (hd=256/512, GQA=8) with multi-step @@ -3384,7 +3416,16 @@ def test_cuda_graph_decode_high_gqa(self): @unittest.mock.patch( "tensorrt_llm.runtime.kv_cache_manager_v2._utils.assert_critical", lambda *a, **kw: None ) - def _run_cuda_graph_real_headdim(self, config_dict, label=""): + def _run_cuda_graph_real_headdim( + self, + config_dict: dict, + label: str = "", + batch_size: int = 2, + initial_cached: list[int] | None = None, + replay_cached: list[int] | None = None, + num_blocks: int = 16, + expect_split_kv: bool = False, + ) -> None: """Helper: CUDA graph decode test with real head_dim configs.""" from tensorrt_llm._torch.attention_backend import ( FlashInferAttention, @@ -3393,17 +3434,29 @@ def _run_cuda_graph_real_headdim(self, config_dict, label=""): from tensorrt_llm._torch.metadata import KVCacheParams config = Gemma4TextConfig(**config_dict) - batch_size = 2 kv_cache_manager = self._get_kv_cache_manager( - config, num_blocks=16, tokens_per_block=32, batch_size=batch_size + config, + num_blocks=num_blocks, + tokens_per_block=32, + batch_size=batch_size, ) self.assertTrue(kv_cache_manager.is_vswa, f"{label}: Expected VSWA manager") request_ids = list(range(batch_size)) - initial_cached = [30, 45] - token_nums = [t + 1 for t in initial_cached] - kv_cache_manager.add_dummy_requests(request_ids, token_nums) + if initial_cached is None: + initial_cached = [30, 45] + self.assertEqual(len(initial_cached), batch_size) + if replay_cached is None: + replay_cached = initial_cached + self.assertEqual(len(replay_cached), batch_size) + reserved_cached = [ + max(initial, replay) + for initial, replay in zip(initial_cached, replay_cached, strict=True) + ] + token_nums = [t + 1 for t in reserved_cached] + requests = kv_cache_manager.add_dummy_requests(request_ids, token_nums, is_gen=True) + self.assertIsNotNone(requests) for i in range(config.num_hidden_layers): buf = kv_cache_manager.get_buffers(i) @@ -3478,21 +3531,6 @@ def _run_cuda_graph_real_headdim(self, config_dict, label=""): ) ) - # --- Reference (eager) --- - ref_metadata = FlashInferAttentionMetadata( - seq_lens=seq_lens, - num_contexts=0, - kv_cache_params=KVCacheParams(use_cache=True, num_cached_tokens_per_seq=initial_cached), - max_num_requests=batch_size, - max_num_tokens=8192, - kv_cache_manager=kv_cache_manager, - request_ids=request_ids, - ) - ref_metadata.prepare() - ref_results = [] - for i in range(num_layers): - ref_results.append(layers[i].forward(gen_qs[i], gen_ks[i], gen_vs[i], ref_metadata)) - # --- CUDA graph --- workspace = torch.empty(320 * 1024 * 1024, dtype=torch.uint8, device="cuda") cg_metadata = FlashInferAttentionMetadata( @@ -3517,21 +3555,62 @@ def _run_cuda_graph_real_headdim(self, config_dict, label=""): with torch.cuda.graph(graph): for i in range(num_layers): cg_results.append(layers[i].forward(gen_qs[i], gen_ks[i], gen_vs[i], cg_metadata)) - graph.replay() - for i in range(num_layers): - torch.testing.assert_close( - cg_results[i], - ref_results[i], - atol=1e-2, - rtol=0, - msg=( - f"{label} Layer {i} ({layer_types[i]}, " - f"hd={layers_info[i]['head_dim']}, " - f"kv={layers_info[i]['num_kv_heads']}): " - f"CUDA graph diverges from eager" + if expect_split_kv: + split_kv_head_dims = set() + for plan_params, wrappers in cg_metadata._plan_params_to_wrappers.items(): + decode_wrapper = wrappers.decode_wrapper + if decode_wrapper is None or decode_wrapper._backend != "fa2": + continue + self.assertTrue( + decode_wrapper._plan_info[-1], + f"{label}: FA2 hd={plan_params.head_dim} did not enable split-K", + ) + split_kv_head_dims.add(plan_params.head_dim) + self.assertEqual(split_kv_head_dims, {256, 512}) + + replay_cached_steps = [replay_cached] + if expect_split_kv: + replay_cached_steps.append([cached + 1 for cached in replay_cached]) + for replay_step, reference_cached in enumerate(replay_cached_steps): + cg_metadata.kv_cache_params = KVCacheParams( + use_cache=True, num_cached_tokens_per_seq=reference_cached + ) + cg_metadata.prepare() + + graph.replay() + + # --- Reference (eager) --- + ref_metadata = FlashInferAttentionMetadata( + seq_lens=seq_lens, + num_contexts=0, + kv_cache_params=KVCacheParams( + use_cache=True, + num_cached_tokens_per_seq=reference_cached, ), + max_num_requests=batch_size, + max_num_tokens=8192, + kv_cache_manager=kv_cache_manager, + request_ids=request_ids, ) + ref_metadata.prepare() + ref_results = [] + for i in range(num_layers): + ref_results.append(layers[i].forward(gen_qs[i], gen_ks[i], gen_vs[i], ref_metadata)) + + for i in range(num_layers): + torch.testing.assert_close( + cg_results[i], + ref_results[i], + atol=1e-2, + rtol=0, + msg=( + f"{label} replay {replay_step}, Layer {i} ({layer_types[i]}, " + f"hd={layers_info[i]['head_dim']}, " + f"kv={layers_info[i]['num_kv_heads']}): " + f"CUDA graph diverges from eager" + ), + ) kv_cache_manager.shutdown() @@ -3543,6 +3622,27 @@ def test_cuda_graph_decode_real_headdim(self): """E2B-like: GQA=8, hd=256/512, non-K=V.""" self._run_cuda_graph_real_headdim(deepcopy(GEMMA4_E2B_REAL_DIMS_CONFIG), "E2B") + @torch.no_grad() + @unittest.mock.patch( + "tensorrt_llm.runtime.kv_cache_manager_v2._utils.assert_critical", lambda *a, **kw: None + ) + @unittest.skipUnless( + torch.cuda.is_available() and torch.cuda.get_device_capability() == (9, 0), + "FA2 split-K schedule refresh is Hopper-specific", + ) + def test_cuda_graph_split_kv_schedule_refresh(self) -> None: + """FA2 split-K graphs refresh schedules for new KV distributions.""" + batch_size = 8 + self._run_cuda_graph_real_headdim( + deepcopy(GEMMA4_12B_REAL_DIMS_CONFIG), + "12B split-K schedule refresh", + batch_size=batch_size, + initial_cached=[4095] + [31] * (batch_size - 1), + replay_cached=[510] * batch_size, + num_blocks=512, + expect_split_kv=True, + ) + @torch.no_grad() @unittest.mock.patch( "tensorrt_llm.runtime.kv_cache_manager_v2._utils.assert_critical", lambda *a, **kw: None @@ -3559,11 +3659,12 @@ def test_cuda_graph_decode_26b_like(self): """26B-like: GQA=2, K=V, hd=256/512.""" self._run_cuda_graph_real_headdim(deepcopy(GEMMA4_26B_REAL_DIMS_CONFIG), "26B") + @unittest.skipUnless(is_sm_100f(), "trtllm-gen attention requires SM100f") @torch.no_grad() @unittest.mock.patch( "tensorrt_llm.runtime.kv_cache_manager_v2._utils.assert_critical", lambda *a, **kw: None ) - def test_cuda_graph_multi_step_trtllm_gen(self): + def test_cuda_graph_multi_step_trtllm_gen(self) -> None: """Multi-step CG decode with trtllm-gen (hd=256/512). Verifies _block_tables update in prepare() works correctly