From 56a1f5ec1e9a75e87d77bd91bbe285519ae97fbb Mon Sep 17 00:00:00 2001
From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
Date: Tue, 28 Apr 2026 20:59:43 +0000
Subject: [PATCH 1/4] [https://nvbugs/5615248][fix] Broader capture of
 piecewise cudagraph

Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
---
 .../_torch/pyexecutor/model_engine.py         |  61 ++++++++-
 tensorrt_llm/llmapi/llm_args.py               |   7 -
 tests/unittest/llmapi/test_llm_args.py        | 126 ++++++++++++++++++
 3 files changed, 183 insertions(+), 11 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index 86911732b02d..a01d373eb09b 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -98,6 +98,47 @@ def warmup(self, resource_manager: ResourceManager) -> None:
         return
 
 
+def _filter_piecewise_capture_num_tokens(
+    candidate_num_tokens: list[int],
+    max_num_tokens: int,
+    max_batch_size: int,
+    max_seq_len: int,
+    num_extra_decoding_steps: int = 0,
+) -> Tuple[list[int], list[int]]:
+    """Cap a piecewise CUDA graph capture-set at the largest forward-pass
+    `num_tokens` value the engine can ever construct.
+
+    A piecewise CUDA graph for `num_tokens = N` is only useful if a real
+    forward pass can ever flow N tokens through the model. Every in-flight
+    request must leave room for at least one decode token in its KV cache,
+    so the largest forward-pass num_tokens reachable at runtime (and the
+    largest the warmup builder can construct) is
+    `max_batch_size * (max_seq_len - 1 - num_extra_decoding_steps)`.
+    Without this cap, candidate entries above this bound would be advertised
+    in `_torch_compile_backend.capture_num_tokens` but silently fail warmup
+    in `_create_warmup_request`, making the outer padding logic pad to a
+    target that has no captured graph and fall back to eager (NVBug 5615248).
+
+    Returns:
+        kept: candidate entries that are <= min(max_num_tokens,
+            max_batch_size * (max_seq_len - 1 - num_extra_decoding_steps)),
+            preserving input order.
+        unrecordable: sorted unique candidate entries that are <=
+            max_num_tokens but exceed the engine's reachable ceiling. These
+            are the entries that would have been silently skipped by warmup.
+    """
+    max_capturable_num_tokens = max(
+        0, max_batch_size * (max_seq_len - 1 - num_extra_decoding_steps))
+    piecewise_capacity_limit = min(max_num_tokens, max_capturable_num_tokens)
+    kept = [i for i in candidate_num_tokens if i <= piecewise_capacity_limit]
+    unrecordable = sorted({
+        i
+        for i in candidate_num_tokens
+        if max_capturable_num_tokens < i <= max_num_tokens
+    })
+    return kept, unrecordable
+
+
 def _filter_cuda_graph_batch_sizes(cuda_graph_batch_sizes: list[int],
                                    max_batch_size: int, max_num_tokens: int,
                                    max_total_draft_tokens: int,
@@ -285,10 +326,22 @@ def __init__(
             torch_compile_piecewise_cuda_graph_num_tokens
             or cuda_graph_batch_sizes or [])
 
-        self._piecewise_cuda_graph_num_tokens = [
-            i for i in piecewise_cuda_graph_num_tokens
-            if i <= self.max_num_tokens
-        ]
+        num_extra_decoding_steps = self._get_num_extra_decoding_steps()
+        self._piecewise_cuda_graph_num_tokens, unrecordable = (
+            _filter_piecewise_capture_num_tokens(
+                piecewise_cuda_graph_num_tokens,
+                max_num_tokens=self.max_num_tokens,
+                max_batch_size=self.batch_size,
+                max_seq_len=self.max_seq_len,
+                num_extra_decoding_steps=num_extra_decoding_steps,
+            ))
+        if unrecordable:
+            logger.warning(
+                f"Skipping piecewise CUDA graph capture for num_tokens="
+                f"{unrecordable}: exceeds reachable ceiling "
+                f"max_batch_size*(max_seq_len-1-num_extra_decoding_steps)="
+                f"{max(0, self.batch_size * (self.max_seq_len - 1 - num_extra_decoding_steps))}. "
+                f"Raise max_seq_len to capture larger graphs.")
 
         try:
             use_ub_for_nccl = (
diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py
index 0a1e7490665a..825a63adb8d8 100644
--- a/tensorrt_llm/llmapi/llm_args.py
+++ b/tensorrt_llm/llmapi/llm_args.py
@@ -3520,13 +3520,6 @@ def validate_capture_num_tokens(cls, v):
         description=
         "The maximum number of CUDA streams to use for torch.compile.")
 
-    @model_validator(mode='after')
-    def set_default_capture_num_tokens(self) -> 'TorchCompileConfig':
-        if self.enable_piecewise_cuda_graph and self.capture_num_tokens is None:
-            self.capture_num_tokens = [2**i for i in range(8)
-                                       ] + [i for i in range(256, 3073, 256)]
-        return self
-
 
 class TorchLlmArgs(BaseLlmArgs):
     # PyTorch backend specific configurations
diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py
index eee83183ced8..67799a46aaa3 100644
--- a/tests/unittest/llmapi/test_llm_args.py
+++ b/tests/unittest/llmapi/test_llm_args.py
@@ -723,6 +723,132 @@ def test_generate_cuda_graph_batch_sizes_padding_edge_cases(
         assert max_batch_size in batch_sizes
 
 
+class TestPiecewiseCudaGraphCaptureDefaults:
+    """Tests for the piecewise CUDA graph capture-set defaults and the
+    model-engine filter that prunes entries the runtime can never reach.
+
+    Two invariants are exercised:
+
+    1. `TorchCompileConfig.capture_num_tokens` defaults to `None`, so the
+       model engine can fall back to `cuda_graph_config.batch_sizes` (the
+       behavior documented on the field).
+    2. `_filter_piecewise_capture_num_tokens` caps the candidate list at
+       `max_batch_size * (max_seq_len - 1 - num_extra_decoding_steps)` --
+       the largest forward-pass `num_tokens` the warmup builder can
+       construct, since every in-flight request must leave room for at
+       least one decode token.
+    """
+
+    def test_torch_compile_config_capture_num_tokens_defaults_to_none(self):
+        """`capture_num_tokens` is unset by default so the documented
+        fallback to `cuda_graph_config.batch_sizes` in the model engine
+        remains reachable. A non-None default here would silently shadow
+        the user's `cuda_graph_config.batch_sizes` choice.
+        """
+        config = TorchCompileConfig(enable_piecewise_cuda_graph=True)
+        assert config.capture_num_tokens is None
+
+    def test_torch_llm_args_capture_num_tokens_defaults_to_none(self):
+        """Same invariant when the config is built through `TorchLlmArgs`
+        (the path real users hit via `trtllm-serve` YAML).
+        """
+        args = TorchLlmArgs(
+            model=llama_model_path,
+            max_batch_size=1,
+            max_seq_len=128,
+            max_beam_width=10,
+            enable_chunked_prefill=True,
+            cuda_graph_config=CudaGraphConfig(max_batch_size=128,
+                                              enable_padding=True),
+            torch_compile_config=TorchCompileConfig(
+                enable_piecewise_cuda_graph=True),
+        )
+        assert args.torch_compile_config.capture_num_tokens is None
+
+    def test_piecewise_filter_drops_entries_above_reachable_ceiling(self):
+        """When the candidate list contains an entry above the reachable
+        ceiling `max_batch_size * (max_seq_len - 1)`, the filter must drop
+        it from `kept` and surface it in `unrecordable`. Otherwise the
+        warmup loop would silently skip it and the outer padding logic
+        would pad to a target with no captured graph.
+        """
+        from tensorrt_llm._torch.pyexecutor.model_engine import (
+            _filter_piecewise_capture_num_tokens)
+
+        candidates = CudaGraphConfig._generate_cuda_graph_batch_sizes(
+            128, enable_padding=True)
+        max_capturable = 1 * (128 - 1)
+        # Precondition: candidate list contains at least one entry above
+        # the reachable ceiling, otherwise the assertions below are vacuous.
+        assert any(i > max_capturable for i in candidates), (
+            "Test precondition no longer holds: cuda_graph_batch_sizes "
+            f"for max_batch_size=128 no longer contains entries above "
+            f"{max_capturable}. Update this test if CudaGraphConfig "
+            "behavior changed.")
+
+        kept, unrecordable = _filter_piecewise_capture_num_tokens(
+            candidates,
+            max_num_tokens=128,
+            max_batch_size=1,
+            max_seq_len=128,
+        )
+
+        assert max(kept) == 120
+        assert 128 not in kept
+        assert unrecordable == [128]
+
+    def test_piecewise_filter_keeps_all_entries_when_within_ceiling(self):
+        """Symmetric case: when the ceiling
+        `max_batch_size * (max_seq_len - 1)` is at least as large as the
+        biggest candidate, nothing is dropped and `unrecordable` is empty.
+        """
+        from tensorrt_llm._torch.pyexecutor.model_engine import (
+            _filter_piecewise_capture_num_tokens)
+
+        candidates = CudaGraphConfig._generate_cuda_graph_batch_sizes(
+            128, enable_padding=True)
+        kept, unrecordable = _filter_piecewise_capture_num_tokens(
+            candidates,
+            max_num_tokens=129,
+            max_batch_size=1,
+            max_seq_len=129,
+        )
+        assert max(kept) == 128
+        assert 128 in kept
+        assert unrecordable == []
+
+    def test_piecewise_filter_subtracts_extra_decoding_steps(self):
+        """Drafting loops consume extra decode steps. The filter must
+        subtract `num_extra_decoding_steps` from the ceiling, mirroring
+        the `max_seq_len - 1 - num_extra_decoding_steps` constraint
+        applied when warmup requests are built.
+        """
+        from tensorrt_llm._torch.pyexecutor.model_engine import (
+            _filter_piecewise_capture_num_tokens)
+
+        candidates = [1, 2, 4, 8, 16, 32, 64, 100, 120]
+        # max_seq_len=128, batch=1, 5 extra decoding steps -> ceiling 122.
+        kept, unrecordable = _filter_piecewise_capture_num_tokens(
+            candidates,
+            max_num_tokens=128,
+            max_batch_size=1,
+            max_seq_len=128,
+            num_extra_decoding_steps=5,
+        )
+        assert max(kept) == 120
+        assert unrecordable == []
+        # Same setup with 9 extra decoding steps -> ceiling 118; 120 drops.
+        kept, unrecordable = _filter_piecewise_capture_num_tokens(
+            candidates,
+            max_num_tokens=128,
+            max_batch_size=1,
+            max_seq_len=128,
+            num_extra_decoding_steps=9,
+        )
+        assert max(kept) == 100
+        assert unrecordable == [120]
+
+
 class TestTrtLlmArgs:
 
     def test_dynamic_setattr(self):

From 3f85137786f0e3f9cbbf307e257364bed535012d Mon Sep 17 00:00:00 2001
From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
Date: Tue, 28 Apr 2026 21:36:22 +0000
Subject: [PATCH 2/4] Clean up

Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
---
 .../_torch/pyexecutor/model_engine.py         | 28 ++------
 tests/unittest/llmapi/test_llm_args.py        | 71 ++++++++++---------
 2 files changed, 43 insertions(+), 56 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index a01d373eb09b..63e6e5268391 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -105,27 +105,13 @@ def _filter_piecewise_capture_num_tokens(
     max_seq_len: int,
     num_extra_decoding_steps: int = 0,
 ) -> Tuple[list[int], list[int]]:
-    """Cap a piecewise CUDA graph capture-set at the largest forward-pass
-    `num_tokens` value the engine can ever construct.
-
-    A piecewise CUDA graph for `num_tokens = N` is only useful if a real
-    forward pass can ever flow N tokens through the model. Every in-flight
-    request must leave room for at least one decode token in its KV cache,
-    so the largest forward-pass num_tokens reachable at runtime (and the
-    largest the warmup builder can construct) is
-    `max_batch_size * (max_seq_len - 1 - num_extra_decoding_steps)`.
-    Without this cap, candidate entries above this bound would be advertised
-    in `_torch_compile_backend.capture_num_tokens` but silently fail warmup
-    in `_create_warmup_request`, making the outer padding logic pad to a
-    target that has no captured graph and fall back to eager (NVBug 5615248).
-
-    Returns:
-        kept: candidate entries that are <= min(max_num_tokens,
-            max_batch_size * (max_seq_len - 1 - num_extra_decoding_steps)),
-            preserving input order.
-        unrecordable: sorted unique candidate entries that are <=
-            max_num_tokens but exceed the engine's reachable ceiling. These
-            are the entries that would have been silently skipped by warmup.
+    """Cap piecewise CUDA graph capture candidates at the engine's reachable
+    `num_tokens` ceiling `max_batch_size * (max_seq_len - 1 - num_extra_decoding_steps)`
+    (each in-flight request must leave room for one decode token).
+
+    Returns `(kept, unrecordable)` where `kept` preserves input order and
+    `unrecordable` is the sorted unique set of entries above the ceiling
+    but within `max_num_tokens` (i.e. would silently fail warmup).
     """
     max_capturable_num_tokens = max(
         0, max_batch_size * (max_seq_len - 1 - num_extra_decoding_steps))
diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py
index 67799a46aaa3..48df4f65f0d8 100644
--- a/tests/unittest/llmapi/test_llm_args.py
+++ b/tests/unittest/llmapi/test_llm_args.py
@@ -724,8 +724,7 @@ def test_generate_cuda_graph_batch_sizes_padding_edge_cases(
 
 
 class TestPiecewiseCudaGraphCaptureDefaults:
-    """Tests for the piecewise CUDA graph capture-set defaults and the
-    model-engine filter that prunes entries the runtime can never reach.
+    """Piecewise CUDA graph capture-set defaults and reachable-ceiling filter.
 
     Two invariants are exercised:
 
@@ -740,17 +739,19 @@ class TestPiecewiseCudaGraphCaptureDefaults:
     """
 
     def test_torch_compile_config_capture_num_tokens_defaults_to_none(self):
-        """`capture_num_tokens` is unset by default so the documented
-        fallback to `cuda_graph_config.batch_sizes` in the model engine
-        remains reachable. A non-None default here would silently shadow
-        the user's `cuda_graph_config.batch_sizes` choice.
+        """`capture_num_tokens` is unset by default.
+
+        Keeps the documented fallback to `cuda_graph_config.batch_sizes`
+        in the model engine reachable; a non-None default here would
+        silently shadow the user's `cuda_graph_config.batch_sizes` choice.
         """
         config = TorchCompileConfig(enable_piecewise_cuda_graph=True)
         assert config.capture_num_tokens is None
 
     def test_torch_llm_args_capture_num_tokens_defaults_to_none(self):
-        """Same invariant when the config is built through `TorchLlmArgs`
-        (the path real users hit via `trtllm-serve` YAML).
+        """Same invariant when reached through `TorchLlmArgs` construction.
+
+        This is the path real users hit via `trtllm-serve` YAML.
         """
         args = TorchLlmArgs(
             model=llama_model_path,
@@ -766,14 +767,15 @@ def test_torch_llm_args_capture_num_tokens_defaults_to_none(self):
         assert args.torch_compile_config.capture_num_tokens is None
 
     def test_piecewise_filter_drops_entries_above_reachable_ceiling(self):
-        """When the candidate list contains an entry above the reachable
-        ceiling `max_batch_size * (max_seq_len - 1)`, the filter must drop
-        it from `kept` and surface it in `unrecordable`. Otherwise the
-        warmup loop would silently skip it and the outer padding logic
-        would pad to a target with no captured graph.
+        """Drop candidates above `max_batch_size * (max_seq_len - 1)`.
+
+        Without the cap, the warmup loop would silently skip these entries
+        and the outer padding logic would pad to a target with no captured
+        graph. They must be removed from `kept` and surfaced in
+        `unrecordable` so the warning fires.
         """
-        from tensorrt_llm._torch.pyexecutor.model_engine import (
-            _filter_piecewise_capture_num_tokens)
+        from tensorrt_llm._torch.pyexecutor.model_engine import \
+            _filter_piecewise_capture_num_tokens
 
         candidates = CudaGraphConfig._generate_cuda_graph_batch_sizes(
             128, enable_padding=True)
@@ -798,12 +800,14 @@ def test_piecewise_filter_drops_entries_above_reachable_ceiling(self):
         assert unrecordable == [128]
 
     def test_piecewise_filter_keeps_all_entries_when_within_ceiling(self):
-        """Symmetric case: when the ceiling
-        `max_batch_size * (max_seq_len - 1)` is at least as large as the
-        biggest candidate, nothing is dropped and `unrecordable` is empty.
+        """Keep all candidates when the largest fits within the ceiling.
+
+        Symmetric case: when `max_batch_size * (max_seq_len - 1)` is at
+        least as large as the biggest candidate, nothing is dropped and
+        `unrecordable` is empty.
         """
-        from tensorrt_llm._torch.pyexecutor.model_engine import (
-            _filter_piecewise_capture_num_tokens)
+        from tensorrt_llm._torch.pyexecutor.model_engine import \
+            _filter_piecewise_capture_num_tokens
 
         candidates = CudaGraphConfig._generate_cuda_graph_batch_sizes(
             128, enable_padding=True)
@@ -818,13 +822,14 @@ def test_piecewise_filter_keeps_all_entries_when_within_ceiling(self):
         assert unrecordable == []
 
     def test_piecewise_filter_subtracts_extra_decoding_steps(self):
-        """Drafting loops consume extra decode steps. The filter must
-        subtract `num_extra_decoding_steps` from the ceiling, mirroring
+        """Subtract `num_extra_decoding_steps` from the ceiling.
+
+        Drafting loops consume extra decode steps; the filter must mirror
         the `max_seq_len - 1 - num_extra_decoding_steps` constraint
         applied when warmup requests are built.
         """
-        from tensorrt_llm._torch.pyexecutor.model_engine import (
-            _filter_piecewise_capture_num_tokens)
+        from tensorrt_llm._torch.pyexecutor.model_engine import \
+            _filter_piecewise_capture_num_tokens
 
         candidates = [1, 2, 4, 8, 16, 32, 64, 100, 120]
         # max_seq_len=128, batch=1, 5 extra decoding steps -> ceiling 122.
@@ -1482,7 +1487,6 @@ def _get_all_llm_args_classes():
 
 def _get_all_pydantic_models_from_llm_args():
     """Get all Pydantic models referenced by BaseLlmArgs and its subclasses."""
-
     visited = set()
     models = []
 
@@ -1535,8 +1539,7 @@ def _get_qualified_name(cls: type) -> str:
 
 
 class TestPydanticBestPractices:
-    """
-    Ensure that the user-facing LlmArgs and its subfields follow Pydantic best practices.
+    """Ensure that the user-facing LlmArgs and its subfields follow Pydantic best practices.
     """
 
     # Fields exempt from Pydantic compatibility checks due to typing limitations or other edge cases.
@@ -1570,8 +1573,7 @@ class TestPydanticBestPractices:
 
     def _is_allowed_type(self, annotation, model_cls: type,
                          field_name: str) -> tuple[bool, str]:
-        """
-        Check if a type annotation is allowed for user-facing config fields.
+        """Check if a type annotation is allowed for user-facing config fields.
 
         Allowed:
         - Pydantic models (must inherit from StrictBaseModel)
@@ -1583,7 +1585,6 @@ def _is_allowed_type(self, annotation, model_cls: type,
 
         Returns (is_allowed, reason) tuple.
         """
-
         # Check if this field is exempt (check class and all parent classes)
         for cls in model_cls.__mro__:
             if field_name in self._COMPATIBILITY_EXEMPT_FIELDS.get(cls, []):
@@ -1669,7 +1670,7 @@ def test_all_fields_have_descriptions(self):
 
         if violations:
             pytest.fail(
-                f"The following fields are missing descriptions:\n" +
+                "The following fields are missing descriptions:\n" +
                 "\n".join(violations) +
                 "\n\nPlease add a description to each by using Field(description=\"...\")."
             )
@@ -1695,7 +1696,7 @@ def test_all_fields_have_allowed_types(self):
 
         if violations:
             pytest.fail(
-                f"The following user-facing fields have types that are not allowed:\n"
+                "The following user-facing fields have types that are not allowed:\n"
                 + "\n".join(violations) +
                 "\n\nPlease use Pydantic-compatible types (primitives, Pydantic models that inherit from StrictBaseModel, "
                 "or other compatible types). If this is intentional, add the field "
@@ -1738,7 +1739,7 @@ def test_no_manual_serialization_or_validation_methods(self):
 
         if violations:
             pytest.fail(
-                f"The following models define forbidden methods:\n" +
+                "The following models define forbidden methods:\n" +
                 "\n".join(violations) +
                 "\n\nPydantic models should follow the recommendations above instead."
             )
@@ -1761,7 +1762,7 @@ def test_no_mutable_default_values(self):
 
         if violations:
             pytest.fail(
-                f"The following fields use mutable default values:\n" +
+                "The following fields use mutable default values:\n" +
                 "\n".join(violations) +
                 "\n\nMutable defaults are shared across instances and cause bugs. "
                 "Use Field(default_factory=...) instead.")
@@ -1788,7 +1789,7 @@ def test_no_custom_init_methods(self):
 
         if violations:
             pytest.fail(
-                f"The following Pydantic models define custom __init__ methods:\n"
+                "The following Pydantic models define custom __init__ methods:\n"
                 + "\n".join(violations) +
                 "\n\nThese should be replaced with alternatives like validators, model_post_init, or classmethods. See this test's docstring for more details."
             )

From cb516b822358dea74a82aa017144ee2aa910962f Mon Sep 17 00:00:00 2001
From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
Date: Tue, 28 Apr 2026 22:09:23 +0000
Subject: [PATCH 3/4] Append ceiling also - if not already present

Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
---
 .../_torch/pyexecutor/model_engine.py         | 28 +++++--
 tests/unittest/llmapi/test_llm_args.py        | 77 +++++++++++++++++--
 2 files changed, 90 insertions(+), 15 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index 63e6e5268391..3ff126ba9e58 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -107,16 +107,29 @@ def _filter_piecewise_capture_num_tokens(
 ) -> Tuple[list[int], list[int]]:
     """Cap piecewise CUDA graph capture candidates at the engine's reachable
     `num_tokens` ceiling `max_batch_size * (max_seq_len - 1 - num_extra_decoding_steps)`
-    (each in-flight request must leave room for one decode token).
-
-    Returns `(kept, unrecordable)` where `kept` preserves input order and
-    `unrecordable` is the sorted unique set of entries above the ceiling
-    but within `max_num_tokens` (i.e. would silently fail warmup).
+    and ensure the ceiling itself is captured.
+
+    Each in-flight request must leave room for at least one decode token,
+    so the ceiling is the largest forward-pass `num_tokens` the warmup
+    builder can construct. Including it in the capture set closes the
+    runtime padding gap between the next-largest candidate and the ceiling
+    (otherwise ISLs in that gap have no graph >= them and fall back to
+    eager).
+
+    Returns `(kept, unrecordable)` where `kept` is sorted ascending,
+    deduped, and contains the ceiling whenever it is positive.
+    `unrecordable` is the sorted unique set of input entries above the
+    ceiling but within `max_num_tokens`.
     """
     max_capturable_num_tokens = max(
         0, max_batch_size * (max_seq_len - 1 - num_extra_decoding_steps))
     piecewise_capacity_limit = min(max_num_tokens, max_capturable_num_tokens)
-    kept = [i for i in candidate_num_tokens if i <= piecewise_capacity_limit]
+    kept = sorted(
+        {i
+         for i in candidate_num_tokens if 0 < i <= piecewise_capacity_limit})
+    if piecewise_capacity_limit > 0 and (not kept or kept[-1]
+                                         < piecewise_capacity_limit):
+        kept.append(piecewise_capacity_limit)
     unrecordable = sorted({
         i
         for i in candidate_num_tokens
@@ -327,7 +340,8 @@ def __init__(
                 f"{unrecordable}: exceeds reachable ceiling "
                 f"max_batch_size*(max_seq_len-1-num_extra_decoding_steps)="
                 f"{max(0, self.batch_size * (self.max_seq_len - 1 - num_extra_decoding_steps))}. "
-                f"Raise max_seq_len to capture larger graphs.")
+                f"Capturing the ceiling itself; raise max_seq_len for larger graphs."
+            )
 
         try:
             use_ub_for_nccl = (
diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py
index 48df4f65f0d8..e029ed2b2d4c 100644
--- a/tests/unittest/llmapi/test_llm_args.py
+++ b/tests/unittest/llmapi/test_llm_args.py
@@ -726,7 +726,7 @@ def test_generate_cuda_graph_batch_sizes_padding_edge_cases(
 class TestPiecewiseCudaGraphCaptureDefaults:
     """Piecewise CUDA graph capture-set defaults and reachable-ceiling filter.
 
-    Two invariants are exercised:
+    Three invariants are exercised:
 
     1. `TorchCompileConfig.capture_num_tokens` defaults to `None`, so the
        model engine can fall back to `cuda_graph_config.batch_sizes` (the
@@ -736,6 +736,10 @@ class TestPiecewiseCudaGraphCaptureDefaults:
        the largest forward-pass `num_tokens` the warmup builder can
        construct, since every in-flight request must leave room for at
        least one decode token.
+    3. The reachable ceiling itself is always present in the returned
+       capture set (when positive), so runtime ISLs in the gap between
+       the next-largest candidate and the ceiling get a graph rather
+       than falling back to eager.
     """
 
     def test_torch_compile_config_capture_num_tokens_defaults_to_none(self):
@@ -772,7 +776,8 @@ def test_piecewise_filter_drops_entries_above_reachable_ceiling(self):
         Without the cap, the warmup loop would silently skip these entries
         and the outer padding logic would pad to a target with no captured
         graph. They must be removed from `kept` and surfaced in
-        `unrecordable` so the warning fires.
+        `unrecordable` so the warning fires. The ceiling itself is then
+        appended so ISLs in the gap still get a graph.
         """
         from tensorrt_llm._torch.pyexecutor.model_engine import \
             _filter_piecewise_capture_num_tokens
@@ -795,7 +800,8 @@ def test_piecewise_filter_drops_entries_above_reachable_ceiling(self):
             max_seq_len=128,
         )
 
-        assert max(kept) == 120
+        assert kept[-1] == max_capturable  # ceiling appended
+        assert 120 in kept  # densely-packed entries below ceiling preserved
         assert 128 not in kept
         assert unrecordable == [128]
 
@@ -804,7 +810,8 @@ def test_piecewise_filter_keeps_all_entries_when_within_ceiling(self):
 
         Symmetric case: when `max_batch_size * (max_seq_len - 1)` is at
         least as large as the biggest candidate, nothing is dropped and
-        `unrecordable` is empty.
+        `unrecordable` is empty. The ceiling (128 here) coincides with the
+        largest candidate so no extra entry is appended.
         """
         from tensorrt_llm._torch.pyexecutor.model_engine import \
             _filter_piecewise_capture_num_tokens
@@ -817,8 +824,10 @@ def test_piecewise_filter_keeps_all_entries_when_within_ceiling(self):
             max_batch_size=1,
             max_seq_len=129,
         )
-        assert max(kept) == 128
+        assert kept[-1] == 128
         assert 128 in kept
+        # Ceiling (128) was already in candidates, must not be duplicated.
+        assert kept.count(128) == 1
         assert unrecordable == []
 
     def test_piecewise_filter_subtracts_extra_decoding_steps(self):
@@ -826,7 +835,9 @@ def test_piecewise_filter_subtracts_extra_decoding_steps(self):
 
         Drafting loops consume extra decode steps; the filter must mirror
         the `max_seq_len - 1 - num_extra_decoding_steps` constraint
-        applied when warmup requests are built.
+        applied when warmup requests are built. The ceiling is appended
+        whenever it is strictly greater than the largest surviving
+        candidate.
         """
         from tensorrt_llm._torch.pyexecutor.model_engine import \
             _filter_piecewise_capture_num_tokens
@@ -840,7 +851,8 @@ def test_piecewise_filter_subtracts_extra_decoding_steps(self):
             max_seq_len=128,
             num_extra_decoding_steps=5,
         )
-        assert max(kept) == 120
+        assert kept[-1] == 122
+        assert 120 in kept
         assert unrecordable == []
         # Same setup with 9 extra decoding steps -> ceiling 118; 120 drops.
         kept, unrecordable = _filter_piecewise_capture_num_tokens(
@@ -850,9 +862,58 @@ def test_piecewise_filter_subtracts_extra_decoding_steps(self):
             max_seq_len=128,
             num_extra_decoding_steps=9,
         )
-        assert max(kept) == 100
+        assert kept[-1] == 118
+        assert 100 in kept
+        assert 120 not in kept
         assert unrecordable == [120]
 
+    def test_piecewise_filter_does_not_double_append_ceiling(self):
+        """Ceiling already present in candidates -> not duplicated."""
+        from tensorrt_llm._torch.pyexecutor.model_engine import \
+            _filter_piecewise_capture_num_tokens
+
+        kept, _ = _filter_piecewise_capture_num_tokens(
+            [1, 64, 128],
+            max_num_tokens=129,
+            max_batch_size=1,
+            max_seq_len=129,
+        )
+        assert kept == [1, 64, 128]
+
+    def test_piecewise_filter_returns_empty_when_ceiling_is_zero(self):
+        """`max_seq_len=1` -> ceiling 0 -> nothing captured.
+
+        With ceiling 0 every positive candidate is unrecordable and the
+        ceiling itself is not appended, so the warning "exceeds reachable
+        ceiling 0; raise max_seq_len" fires for the full candidate list.
+        """
+        from tensorrt_llm._torch.pyexecutor.model_engine import \
+            _filter_piecewise_capture_num_tokens
+
+        kept, unrecordable = _filter_piecewise_capture_num_tokens(
+            [1, 2, 4],
+            max_num_tokens=8,
+            max_batch_size=1,
+            max_seq_len=1,
+        )
+        assert kept == []
+        assert unrecordable == [1, 2, 4]
+
+    def test_piecewise_filter_appends_ceiling_when_only_smaller_candidates(
+            self):
+        """No candidate near the ceiling -> ceiling still appended."""
+        from tensorrt_llm._torch.pyexecutor.model_engine import \
+            _filter_piecewise_capture_num_tokens
+
+        kept, _ = _filter_piecewise_capture_num_tokens(
+            [1, 2, 4, 8],
+            max_num_tokens=1024,
+            max_batch_size=8,
+            max_seq_len=128,
+        )
+        # Ceiling: 8 * (128 - 1) = 1016.
+        assert kept == [1, 2, 4, 8, 1016]
+
 
 class TestTrtLlmArgs:
 

From 20cd33f7341c01b6b24d2faa92e0f663bafdbfa1 Mon Sep 17 00:00:00 2001
From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
Date: Mon, 4 May 2026 23:56:29 +0000
Subject: [PATCH 4/4] bring back validator - make sure ceiling is always
 present

Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
---
 tensorrt_llm/llmapi/llm_args.py        |  7 ++++
 tests/unittest/llmapi/test_llm_args.py | 52 ++++++++++++++++++++------
 2 files changed, 48 insertions(+), 11 deletions(-)

diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py
index 825a63adb8d8..0a1e7490665a 100644
--- a/tensorrt_llm/llmapi/llm_args.py
+++ b/tensorrt_llm/llmapi/llm_args.py
@@ -3520,6 +3520,13 @@ def validate_capture_num_tokens(cls, v):
         description=
         "The maximum number of CUDA streams to use for torch.compile.")
 
+    @model_validator(mode='after')
+    def set_default_capture_num_tokens(self) -> 'TorchCompileConfig':
+        if self.enable_piecewise_cuda_graph and self.capture_num_tokens is None:
+            self.capture_num_tokens = [2**i for i in range(8)
+                                       ] + [i for i in range(256, 3073, 256)]
+        return self
+
 
 class TorchLlmArgs(BaseLlmArgs):
     # PyTorch backend specific configurations
diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py
index e029ed2b2d4c..f05c90bcdb99 100644
--- a/tests/unittest/llmapi/test_llm_args.py
+++ b/tests/unittest/llmapi/test_llm_args.py
@@ -728,9 +728,13 @@ class TestPiecewiseCudaGraphCaptureDefaults:
 
     Three invariants are exercised:
 
-    1. `TorchCompileConfig.capture_num_tokens` defaults to `None`, so the
-       model engine can fall back to `cuda_graph_config.batch_sizes` (the
-       behavior documented on the field).
+    1. `TorchCompileConfig.capture_num_tokens` defaults to a fixed
+       powers-of-2 + 256-stride list when `enable_piecewise_cuda_graph`
+       is True (and stays `None` otherwise). The fixed list keeps the
+       capture set small to bound startup time and CUDA graph memory;
+       the model-engine filter (invariants 2 and 3) ensures the largest
+       reachable size is always captured even when it is not in this
+       default list.
     2. `_filter_piecewise_capture_num_tokens` caps the candidate list at
        `max_batch_size * (max_seq_len - 1 - num_extra_decoding_steps)` --
        the largest forward-pass `num_tokens` the warmup builder can
@@ -742,18 +746,44 @@ class TestPiecewiseCudaGraphCaptureDefaults:
        than falling back to eager.
     """
 
-    def test_torch_compile_config_capture_num_tokens_defaults_to_none(self):
-        """`capture_num_tokens` is unset by default.
+    _EXPECTED_DEFAULT_CAPTURE_NUM_TOKENS = [2**i for i in range(8)] + list(
+        range(256, 3073, 256))
 
-        Keeps the documented fallback to `cuda_graph_config.batch_sizes`
-        in the model engine reachable; a non-None default here would
-        silently shadow the user's `cuda_graph_config.batch_sizes` choice.
+    def test_torch_compile_config_capture_num_tokens_default_when_piecewise_enabled(
+            self):
+        """Default capture set is the powers-of-2 + 256-stride list.
+
+        Keeps the capture set bounded (~20 entries) so server startup
+        time and CUDA graph memory stay predictable. The model engine
+        further filters and appends the reachable ceiling, so
+        out-of-range entries (e.g. > max_seq_len-1) are never recorded
+        and gap ISLs still get a graph.
         """
         config = TorchCompileConfig(enable_piecewise_cuda_graph=True)
+        assert config.capture_num_tokens == self._EXPECTED_DEFAULT_CAPTURE_NUM_TOKENS
+
+    def test_torch_compile_config_capture_num_tokens_stays_none_when_piecewise_disabled(
+            self):
+        """No default is populated when piecewise is off.
+
+        The capture set is irrelevant in this case; populating it would
+        be misleading in serialized configs.
+        """
+        config = TorchCompileConfig(enable_piecewise_cuda_graph=False)
         assert config.capture_num_tokens is None
 
-    def test_torch_llm_args_capture_num_tokens_defaults_to_none(self):
-        """Same invariant when reached through `TorchLlmArgs` construction.
+    def test_torch_compile_config_capture_num_tokens_user_override_preserved(
+            self):
+        """User-supplied `capture_num_tokens` is not overwritten by the default."""
+        user_list = [4, 8, 16]
+        config = TorchCompileConfig(enable_piecewise_cuda_graph=True,
+                                    capture_num_tokens=user_list)
+        # `validate_capture_num_tokens` dedupes and reverse-sorts.
+        assert config.capture_num_tokens == sorted(set(user_list), reverse=True)
+
+    def test_torch_llm_args_capture_num_tokens_default_when_piecewise_enabled(
+            self):
+        """Same default applies when reached through `TorchLlmArgs` construction.
 
         This is the path real users hit via `trtllm-serve` YAML.
         """
@@ -768,7 +798,7 @@ def test_torch_llm_args_capture_num_tokens_defaults_to_none(self):
             torch_compile_config=TorchCompileConfig(
                 enable_piecewise_cuda_graph=True),
         )
-        assert args.torch_compile_config.capture_num_tokens is None
+        assert args.torch_compile_config.capture_num_tokens == self._EXPECTED_DEFAULT_CAPTURE_NUM_TOKENS
 
     def test_piecewise_filter_drops_entries_above_reachable_ceiling(self):
         """Drop candidates above `max_batch_size * (max_seq_len - 1)`.