diff --git a/jenkins/L0_Test.groovy b/jenkins/L0_Test.groovy index ee9654166f9c..8f2d23bd26a9 100644 --- a/jenkins/L0_Test.groovy +++ b/jenkins/L0_Test.groovy @@ -4957,6 +4957,8 @@ def launchTestJobs(pipeline, testFilter) "DGX_B300-4_GPUs-PyTorch-Post-Merge-2": ["auto:dgx-b300-flex", "l0_dgx_b300", 2, 2, 4, 1, true], // VisualGen PerfSanity post-merge test "DGX_B200-8_GPUs-PyTorch-VisualGen-PerfSanity-Post-Merge-1": ["auto:dgx-b200-flex", "l0_b200_visual_gen_perf_sanity", 1, 1, 8, 1, true], + // Single-GPU Gemma4 PerfSanity regression gate and baseline + "DGX_B200-PyTorch-PerfSanity-1": ["auto:dgx-b200-flex", "l0_b200_perf_sanity", 1, 1, 1, 1, true], // PerfSanity post-merge tests "DGX_B200-8_GPUs-PyTorch-PerfSanity-Post-Merge-1": ["auto:dgx-b200-flex", "l0_b200_multi_gpus_perf_sanity", 1, 4, 8, 1, true], "DGX_B200-8_GPUs-PyTorch-PerfSanity-Post-Merge-2": ["auto:dgx-b200-flex", "l0_b200_multi_gpus_perf_sanity", 2, 4, 8, 1, true], diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py index 6eb9c366ff7c..933c1dc443b9 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py @@ -175,6 +175,14 @@ class DeepseekV4CacheManager(KVCacheManagerV2): # For other attention types, block size is tokens_per_block. compressed_block_sizes: List[int] + def _get_typical_seq_len(self, kv_cache_config: KvCacheConfig) -> int: + """Retain DeepSeek-V4's max-length pool-sizing model by default.""" + return ( + kv_cache_config.avg_seq_len + if kv_cache_config.avg_seq_len is not None + else self.max_seq_len + ) + def __init__( self, kv_cache_config: KvCacheConfig, diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 24808e8b55b5..e9e6b14135fc 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -1658,59 +1658,61 @@ def _build_base_config( typical_step = None constraints = [] if kv_cache_config.pool_ratio is None: - typical_seq_len = ( - kv_cache_config.avg_seq_len - if kv_cache_config.avg_seq_len is not None - else self.max_seq_len - ) - if typical_seq_len > self.max_seq_len: + typical_seq_len = self._get_typical_seq_len(kv_cache_config) + if typical_seq_len is not None and typical_seq_len > self.max_seq_len: raise ValueError( f"kv_cache_config.avg_seq_len ({typical_seq_len}) must be less than or " f"equal to max_seq_len ({self.max_seq_len})" ) - # Model one context request and enough generation requests to fill - # max_batch_size without over-provisioning windowed cache pools. - context_capacity = ( - self.max_num_tokens if self.max_num_tokens is not None else typical_seq_len - ) + self.num_extra_kv_tokens - generation_history_length = max(0, typical_seq_len - self.max_draft_len - 1) - typical_step = BatchDesc( - [KVCacheDesc(capacity=context_capacity, history_length=0)] - + [ - KVCacheDesc( - capacity=typical_seq_len, - history_length=generation_history_length, - ) - ] - * (self.max_batch_size - 1) - ) - - # CUDA graph generation warmup uses one request at max_seq_len and - # enough minimal decode requests to fill max_batch_size. - min_decode_capacity = 1 + self.max_draft_len + self.num_extra_kv_tokens - constraints.append( - BatchDesc( - [KVCacheDesc(capacity=self.max_seq_len, history_length=self.max_seq_len - 1)] - + [KVCacheDesc(capacity=min_decode_capacity, history_length=0)] + if typical_seq_len is not None: + # Model one context request and enough generation requests to fill + # max_batch_size without over-provisioning windowed cache pools. + context_capacity = ( + self.max_num_tokens if self.max_num_tokens is not None else typical_seq_len + ) + self.num_extra_kv_tokens + generation_history_length = max(0, typical_seq_len - self.max_draft_len - 1) + typical_step = BatchDesc( + [KVCacheDesc(capacity=context_capacity, history_length=0)] + + [ + KVCacheDesc( + capacity=typical_seq_len, + history_length=generation_history_length, + ) + ] * (self.max_batch_size - 1) ) - ) - # General and chunked-prefill warmup uses one fresh context request - # at the per-iteration token budget. - if self.max_num_tokens is not None: + # CUDA graph generation warmup uses one request at max_seq_len and + # enough minimal decode requests to fill max_batch_size. + min_decode_capacity = 1 + self.max_draft_len + self.num_extra_kv_tokens constraints.append( BatchDesc( [ KVCacheDesc( - capacity=self.max_num_tokens + self.num_extra_kv_tokens, - history_length=0, + capacity=self.max_seq_len, + history_length=self.max_seq_len - 1, ) ] + + [KVCacheDesc(capacity=min_decode_capacity, history_length=0)] + * (self.max_batch_size - 1) ) ) + # General and chunked-prefill warmup uses one fresh context request + # at the per-iteration token budget. + if self.max_num_tokens is not None: + constraints.append( + BatchDesc( + [ + KVCacheDesc( + capacity=self.max_num_tokens + self.num_extra_kv_tokens, + history_length=0, + ) + ] + ) + ) + buffer_type = [Role.KEY] if self.kv_cache_type != CacheTypeCpp.SELFKONLY: buffer_type.append(Role.VALUE) @@ -1782,6 +1784,10 @@ def _build_cache_config(self, config: KVCacheManagerConfigPy) -> KVCacheManagerC """Customize the general cache config for a specialized cache manager.""" return config + def _get_typical_seq_len(self, kv_cache_config: KvCacheConfig) -> int | None: + """Return the configured typical sequence length, if any.""" + return kv_cache_config.avg_seq_len + def _extra_buffers_per_layer( self, *, tokens_per_block: int ) -> Optional[dict[int, List[BufferConfig]]]: diff --git a/tests/integration/defs/perf/_model_paths.py b/tests/integration/defs/perf/_model_paths.py index 15ebbe1d0380..c909a85ea3f6 100644 --- a/tests/integration/defs/perf/_model_paths.py +++ b/tests/integration/defs/perf/_model_paths.py @@ -45,6 +45,7 @@ "gemma_3_12b_it_fp8": "gemma/gemma-3-12b-it-fp8", "gemma_3_12b_it_fp4": "gemma/gemma-3-12b-it-fp4", "gemma_3_1b_it": "gemma/gemma-3-1b-it", + "gemma_4_26b_a4b_nvfp4": "gemma/nvidia-Gemma-4-26B-A4B-NVFP4", "deepseek_r1_fp8": "DeepSeek-R1/DeepSeek-R1", "deepseek_r1_nvfp4": "DeepSeek-R1/DeepSeek-R1-FP4", "deepseek_r1_0528_fp8": "DeepSeek-R1/DeepSeek-R1-0528/", diff --git a/tests/integration/test_lists/test-db/l0_b200_perf_sanity.yml b/tests/integration/test_lists/test-db/l0_b200_perf_sanity.yml new file mode 100644 index 000000000000..3ef357924a0a --- /dev/null +++ b/tests/integration/test_lists/test-db/l0_b200_perf_sanity.yml @@ -0,0 +1,22 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +version: 0.0.1 +l0_b200_perf_sanity: +- condition: + ranges: + system_gpu_count: + gte: 1 + lte: 1 + wildcards: + gpu: + - '*b200*' + linux_distribution_name: ubuntu* + cpu: x86_64 + terms: + stage: pre_merge + backend: pytorch + orchestrator: mpi + tests: + # gemma-4-26b-a4b-nvfp4 + - perf/test_perf_sanity.py::test_e2e[aggr_upload-gemma4_26b_a4b_nvfp4_blackwell-gemma4_26b_a4b_nvfp4_tp1_1k1k] TIMEOUT (60) diff --git a/tests/scripts/perf-sanity/aggregated/gemma4_26b_a4b_nvfp4_blackwell.yaml b/tests/scripts/perf-sanity/aggregated/gemma4_26b_a4b_nvfp4_blackwell.yaml new file mode 100644 index 000000000000..a9997d646100 --- /dev/null +++ b/tests/scripts/perf-sanity/aggregated/gemma4_26b_a4b_nvfp4_blackwell.yaml @@ -0,0 +1,33 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +metadata: + model_name: gemma_4_26b_a4b_nvfp4 + supported_gpus: + - B200 +hardware: + gpus_per_node: 1 +server_configs: + # Production-like concurrency 64 covers the scheduler, KV-pool sizing, + # CUDA graph padding, and VSWA block-table decode paths together. + - name: "gemma4_26b_a4b_nvfp4_tp1_1k1k" + tensor_parallel_size: 1 + max_batch_size: 256 + load_format: dummy + trust_remote_code: true + enable_chunked_prefill: true + cuda_graph_config: + enable_padding: true + max_batch_size: 256 + kv_cache_config: + dtype: fp8 + enable_block_reuse: false + multimodal_config: + encoder_cache_max_bytes: 0 + client_configs: + - name: "con64_iter3_1k1k" + concurrency: 64 + iterations: 3 + isl: 1000 + osl: 1000 + backend: "openai-chat" diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py index bcba3baf1b24..8161ac5f3525 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py @@ -50,6 +50,16 @@ ] +@pytest.mark.parametrize(("avg_seq_len", "expected"), [(None, 1024), (256, 256)]) +def test_typical_seq_len_preserves_deepseek_v4_fallback( + avg_seq_len: int | None, expected: int +) -> None: + manager = object.__new__(DeepseekV4CacheManager) + manager.max_seq_len = 1024 + + assert manager._get_typical_seq_len(KvCacheConfig(avg_seq_len=avg_seq_len)) == expected + + def test_cache_size_estimation_uses_model_attention_layer_count(): class FakeModelConfig: sparse_attention_config = SimpleNamespace( diff --git a/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py b/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py index e5245db947fc..399fe62b1cc8 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py +++ b/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py @@ -136,7 +136,7 @@ def test_pool_ratio_overrides_constraints() -> None: assert config.constraints == [] -def test_builds_warmup_constraints() -> None: +def test_default_uses_allocator_fallback() -> None: config = _make_cache_config_for_test( KvCacheConfig(host_cache_size=0), max_batch_size=3, @@ -146,6 +146,19 @@ def test_builds_warmup_constraints() -> None: ) assert config.initial_pool_ratio is None + assert config.typical_step is None + assert config.constraints == [] + + +def test_avg_seq_len_builds_warmup_constraints() -> None: + config = _make_cache_config_for_test( + KvCacheConfig(host_cache_size=0, avg_seq_len=1024), + max_batch_size=3, + max_seq_len=1024, + max_num_tokens=2048, + max_draft_len=2, + ) + assert config.typical_step == BatchDesc( [KVCacheDesc(capacity=2048, history_length=0)] + [KVCacheDesc(capacity=1024, history_length=1021)] * 2 @@ -187,7 +200,7 @@ def test_avg_seq_len_must_not_exceed_max_seq_len() -> None: def test_extra_tokens_are_in_context_capacity() -> None: config = _make_cache_config_for_test( - KvCacheConfig(), + KvCacheConfig(avg_seq_len=264), max_batch_size=1, max_seq_len=264, max_num_tokens=256,