[None][feat] Support ignored prompt length for penalties via new sampling config parameter #8127

nvxuanyuc · 2025-10-03T00:12:39Z

… to ignore from the prompt for precense and frequency penalties for trt and torch path.

Summary by CodeRabbit

New Features
- Added prompt_ignore_length to generation, allowing part of the prompt to be excluded from repetition/presence/frequency penalties. Supports scalar or per-batch values.
- Available via Python SamplingParams, generate() kwargs, dynamic decode setup, and OpenAI-compatible Completion/ChatCompletion requests.
- New CLI flag --prompt_ignore_length.
Examples
- Updated example scripts to accept and pass prompt_ignore_length through the generation pipeline (including long-context eval, run, and summarize workflows).

Description

Currently, penalties are always applied across both prompt and generated tokens, with no way to restrict them to only generated tokens. This PR adds a new parameter prompt_ignore_length to sampling config, allowing dynamic control over how many prompt tokens are excluded, making the sampling penalty logic more flexible.

Test Coverage

Add additional sampling kernel unit tests, including:

Default mode (no prompt tokens ignored)
Partial prompt ignored
Full prompt ignored

GitHub Bot Help

/bot [-h] ['run', 'kill', 'skip', 'reuse-pipeline'] ...

Provide a user friendly way for developers to interact with a Jenkins server.

Run /bot [-h|--help] to print this help message.

See details below for each supported subcommand.

run [--reuse-test (optional)pipeline-id --disable-fail-fast --skip-test --stage-list "A10-PyTorch-1, xxx" --gpu-type "A30, H100_PCIe" --test-backend "pytorch, cpp" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" --detailed-log --debug(experimental)]

Launch build/test pipelines. All previously running jobs will be killed.

--reuse-test (optional)pipeline-id (OPTIONAL) : Allow the new pipeline to reuse build artifacts and skip successful test stages from a specified pipeline or the last pipeline if no pipeline-id is indicated. If the Git commit ID has changed, this option will be always ignored. The DEFAULT behavior of the bot is to reuse build artifacts and successful test results from the last pipeline.

--disable-reuse-test (OPTIONAL) : Explicitly prevent the pipeline from reusing build artifacts and skipping successful test stages from a previous pipeline. Ensure that all builds and tests are run regardless of previous successes.

--disable-fail-fast (OPTIONAL) : Disable fail fast on build/tests/infra failures.

--skip-test (OPTIONAL) : Skip all test stages, but still run build stages, package stages and sanity check stages. Note: Does NOT update GitHub check status.

--stage-list "A10-PyTorch-1, xxx" (OPTIONAL) : Only run the specified test stages. Examples: "A10-PyTorch-1, xxx". Note: Does NOT update GitHub check status.

--gpu-type "A30, H100_PCIe" (OPTIONAL) : Only run the test stages on the specified GPU types. Examples: "A30, H100_PCIe". Note: Does NOT update GitHub check status.

--test-backend "pytorch, cpp" (OPTIONAL) : Skip test stages which don't match the specified backends. Only support [pytorch, cpp, tensorrt, triton]. Examples: "pytorch, cpp" (does not run test stages with tensorrt or triton backend). Note: Does NOT update GitHub pipeline status.

--only-multi-gpu-test (OPTIONAL) : Only run the multi-GPU tests. Note: Does NOT update GitHub check status.

--disable-multi-gpu-test (OPTIONAL) : Disable the multi-GPU tests. Note: Does NOT update GitHub check status.

--add-multi-gpu-test (OPTIONAL) : Force run the multi-GPU tests in addition to running L0 pre-merge pipeline.

--post-merge (OPTIONAL) : Run the L0 post-merge pipeline instead of the ordinary L0 pre-merge pipeline.

--extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" (OPTIONAL) : Run the ordinary L0 pre-merge pipeline and specified test stages. Examples: --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx".

--detailed-log (OPTIONAL) : Enable flushing out all logs to the Jenkins console. This will significantly increase the log volume and may slow down the job.

--debug (OPTIONAL) : Experimental feature. Enable access to the CI container for debugging purpose. Note: Specify exactly one stage in the stage-list parameter to access the appropriate container environment. Note: Does NOT update GitHub check status.

For guidance on mapping tests to stage names, see docs/source/reference/ci-overview.md
and the scripts/test_to_stage_mapping.py helper.

kill

kill

Kill all running builds associated with pull request.

skip

skip --comment COMMENT

Skip testing for latest commit on pull request. --comment "Reason for skipping build/test" is required. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

reuse-pipeline

reuse-pipeline

Reuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

jhaotingc · 2025-10-03T00:13:50Z

/bot run --disable-fail-fast

coderabbitai · 2025-10-03T00:22:06Z

📝 Walkthrough

Walkthrough

Adds a new optional parameter “promptIgnoreLength” across the stack to control how many initial tokens are ignored by presence/frequency penalties. The change threads this field from public APIs (Python/OpenAI protocol) through bindings, serialization, executor/runtime configs, decoding layers, and into CUDA penalty kernels, with tests and examples updated accordingly.

Changes

Cohort / File(s)	Summary
Executor SamplingConfig API `cpp/include/tensorrt_llm/executor/executor.h`, `cpp/tensorrt_llm/executor/samplingConfig.cpp`, `cpp/tensorrt_llm/executor/serialization.cpp`	Introduces SamplingConfig::promptIgnoreLength (optional SizeType32): ctor parameter, private member, getter/setter, static checker, equality update; serialization/deserialization and size computations updated.
Runtime Sampling Config & Decoder `cpp/include/tensorrt_llm/runtime/samplingConfig.h`, `cpp/tensorrt_llm/runtime/gptDecoder.cpp`	Adds OptVec promptIgnoreLength to runtime::SamplingConfig; propagates to PenaltySetupParams during decoder setup and disable-lookahead path.
Kernels: Penalty `cpp/tensorrt_llm/kernels/penaltyKernels.cu`, `cpp/tensorrt_llm/kernels/penaltyKernels.h`, `cpp/tensorrt_llm/kernels/penaltyTypes.h`	batchApplyPenalty gains promptIgnoreLengths input; workspace expanded to 2×vocabSize; logic updated to ignore first N tokens per-beam for presence/frequency; new DecodingPenaltyType::PromptIgnoreLength and limits handling; launcher params extended.
Layers: Penalty setup/buffers `cpp/tensorrt_llm/layers/penaltyLayer.cpp`, `cpp/tensorrt_llm/layers/penaltyLayer.h`, `cpp/tensorrt_llm/layers/decodingParams.h`, `cpp/include/tensorrt_llm/layers/defaultDecodingParams.h`	Adds PenaltySetupParams::promptIgnoreLength; allocates/handles host/device buffers; doubles penalty workspace when used; forwards promptIgnoreLengths to kernels; DefaultDecodingParams exposes getPromptIgnoreLength() constexpr=0.
Torch Dynamic Decode op `cpp/tensorrt_llm/thop/dynamicDecodeOp.h`, `cpp/tensorrt_llm/thop/dynamicDecodeOp.cpp`	Adds optional prompt_ignore_length tensor parameter to setup signatures (IFtDynamicDecode, FtDynamicDecode, DynamicDecodeOp); threads value into penalty params.
Python runtime `tensorrt_llm/runtime/generation.py`, `tensorrt_llm/runtime/model_runner_cpp.py`	SamplingConfig gains prompt_ignore_length; decoder setup materializes per-batch int32 tensor and passes it to dynamic decoder; generate() kwargs accept prompt_ignore_length.
Bindings (C++ ⇄ Python) `cpp/tensorrt_llm/nanobind/bindings.cpp`, `cpp/tensorrt_llm/nanobind/executor/request.cpp`, `cpp/tensorrt_llm/pybind/bindings.cpp`, `cpp/tensorrt_llm/pybind/executor/request.cpp`	Expose SamplingConfig.prompt_ignore_length property; update getstate/setstate tuple size to 20; ctor/binding arg lists and indices adjusted.
Serving API & Params `tensorrt_llm/serve/openai_protocol.py`, `tensorrt_llm/sampling_params.py`	Adds prompt_ignore_length to CompletionRequest/ChatCompletionRequest and passes through to SamplingParams; SamplingParams includes optional prompt_ignore_length.
Examples `examples/utils.py`, `examples/run.py`, `examples/eval_long_context.py`, `examples/summarize.py`	CLI adds --prompt_ignore_length; generate() calls forward prompt_ignore_length through runners in examples and eval paths.
Tests `cpp/tests/unit_tests/executor/samplingConfigTest.cpp`, `cpp/tests/unit_tests/runtime/samplingConfigTest.cpp`, `cpp/tests/unit_tests/kernels/sampling/samplingPenaltyTest.cpp`, `cpp/tests/unit_tests/layers/dynamicDecodeLayerTest.h`	Update constructors and expectations for new field; extend penalty tests to cover promptIgnoreLengths behavior; expand penalty workspace to 2×vocab; add test params vector for promptIgnoreLengths.

Sequence Diagram(s)

sequenceDiagram
  autonumber
  participant Client
  participant Serve as OpenAI Protocol
  participant Py as Python Runtime
  participant Bind as C++ Bindings
  participant Exec as Executor::SamplingConfig
  participant Run as Runtime::SamplingConfig
  participant Layer as PenaltyLayer
  participant Kern as CUDA Penalty Kernel

  Client->>Serve: Request(prompt_ignore_length=?)
  Serve->>Py: SamplingParams(prompt_ignore_length)
  Py->>Py: Build SamplingConfig(prompt_ignore_length)
  Py->>Bind: set SamplingConfig.prompt_ignore_length
  Bind->>Exec: construct SamplingConfig(..., promptIgnoreLength)
  Exec->>Run: to runtime::SamplingConfig(promptIgnoreLength per-batch)
  Run->>Layer: PenaltySetupParams.promptIgnoreLength
  Layer->>Kern: batchApplyPenalty(..., promptIgnoreLengths, workspace=2×vocab)
  Note over Kern: Ignore first N tokens for presence/frequency

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Pre-merge checks and finishing touches

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Description Check	⚠️ Warning	The PR description includes a Description and Test Coverage section but still contains unremoved template comments, lacks a filled summary in place of the stub “@coderabbitai summary,” and omits the required PR Checklist section.	Remove the commented template instructions, provide a concise summary after the “@coderabbitai summary” marker, and add a PR Checklist section that follows the repository template.
Docstring Coverage	⚠️ Warning	Docstring coverage is 6.52% which is insufficient. The required threshold is 80.00%.	You can run `@coderabbitai generate docstrings` to improve docstring coverage.

✅ Passed checks (1 passed)

Check name	Status	Explanation
Title Check	✅ Passed	The title follows the required “[ticket][type] Summary” format with a valid “None” ticket and “[feat]” type, and it clearly and concisely describes the primary change of adding support for ignored prompt length in the sampling configuration.

✨ Finishing touches

📝 Generate Docstrings

🧪 Generate unit tests

Create PR with unit tests
Post copyable unit tests in a comment

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

coderabbitai

Actionable comments posted: 1

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

cpp/tensorrt_llm/nanobind/executor/request.cpp (1)

84-108: Restore backward compatibility for pickled SamplingConfig state

Older pickles still store 19 fields. With this change __setstate__ now throws for those tuples, so any previously serialized executor/request state can no longer be loaded. Please accept both versions and default promptIgnoreLength when the legacy layout is encountered.

Here is one way to keep both formats working:

-        if (state.size() != 20)
+        if (state.size() != 19 && state.size() != 20)
         {
             throw std::runtime_error("Invalid SamplingConfig state!");
         }
-        new (&samplingConfig) tle::SamplingConfig(nb::cast<SizeType32>(state[0]), // BeamWidth
+        if (state.size() == 20)
+        {
+        new (&samplingConfig) tle::SamplingConfig(nb::cast<SizeType32>(state[0]), // BeamWidth
             …
-            nb::cast<std::optional<FloatType>>(state[18]),                        // MinP
-            nb::cast<std::optional<std::vector<SizeType32>>>(state[19])           // BeamWidthArray
-        );
+            nb::cast<std::optional<FloatType>>(state[18]),                        // MinP
+            nb::cast<std::optional<std::vector<SizeType32>>>(state[19])           // BeamWidthArray
+        );
+        }
+        else
+        {
+        new (&samplingConfig) tle::SamplingConfig(nb::cast<SizeType32>(state[0]), // BeamWidth
+            nb::cast<std::optional<SizeType32>>(state[1]),                        // TopK
+            …
+            nb::cast<std::optional<FloatType>>(state[12]),                        // FrequencyPenalty
+            std::nullopt,                                                        // PromptIgnoreLength (legacy)
+            nb::cast<std::optional<FloatType>>(state[13]),                        // LengthPenalty
+            nb::cast<std::optional<SizeType32>>(state[14]),                       // EarlyStopping
+            nb::cast<std::optional<SizeType32>>(state[15]),                       // NoRepeatNgramSize
+            nb::cast<std::optional<SizeType32>>(state[16]),                       // NumReturnSequences
+            nb::cast<std::optional<FloatType>>(state[17]),                        // MinP
+            nb::cast<std::optional<std::vector<SizeType32>>>(state[18]));         // BeamWidthArray
+        }

🧹 Nitpick comments (2)

examples/utils.py (1)

307-307: Document the new CLI flag.

Please add a short help= string so users know this flag controls how many prompt tokens are exempt from repetition penalties. It keeps the CLI self-documenting alongside the rest of the sampling knobs.

tensorrt_llm/serve/openai_protocol.py (1)

227-228: Clamp invalid prompt_ignore_length early.

Right now callers could send a negative value, which would travel through SamplingParams and eventually the CUDA penalties path without any guard. Please add Field(ge=0) (or similar validation) so we reject bad inputs at the API layer before they hit runtime math.

📜 Review details

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between a5b59fd and 4d33465.

📒 Files selected for processing (30)

cpp/include/tensorrt_llm/executor/executor.h (5 hunks)
cpp/include/tensorrt_llm/layers/defaultDecodingParams.h (1 hunks)
cpp/include/tensorrt_llm/runtime/samplingConfig.h (4 hunks)
cpp/tensorrt_llm/executor/samplingConfig.cpp (6 hunks)
cpp/tensorrt_llm/executor/serialization.cpp (4 hunks)
cpp/tensorrt_llm/kernels/penaltyKernels.cu (7 hunks)
cpp/tensorrt_llm/kernels/penaltyKernels.h (1 hunks)
cpp/tensorrt_llm/kernels/penaltyTypes.h (2 hunks)
cpp/tensorrt_llm/layers/decodingParams.h (1 hunks)
cpp/tensorrt_llm/layers/penaltyLayer.cpp (7 hunks)
cpp/tensorrt_llm/layers/penaltyLayer.h (1 hunks)
cpp/tensorrt_llm/nanobind/bindings.cpp (3 hunks)
cpp/tensorrt_llm/nanobind/executor/request.cpp (5 hunks)
cpp/tensorrt_llm/pybind/bindings.cpp (3 hunks)
cpp/tensorrt_llm/pybind/executor/request.cpp (5 hunks)
cpp/tensorrt_llm/runtime/gptDecoder.cpp (2 hunks)
cpp/tensorrt_llm/thop/dynamicDecodeOp.cpp (5 hunks)
cpp/tensorrt_llm/thop/dynamicDecodeOp.h (3 hunks)
cpp/tests/unit_tests/executor/samplingConfigTest.cpp (4 hunks)
cpp/tests/unit_tests/kernels/sampling/samplingPenaltyTest.cpp (40 hunks)
cpp/tests/unit_tests/layers/dynamicDecodeLayerTest.h (1 hunks)
cpp/tests/unit_tests/runtime/samplingConfigTest.cpp (5 hunks)
examples/eval_long_context.py (1 hunks)
examples/run.py (3 hunks)
examples/summarize.py (2 hunks)
examples/utils.py (1 hunks)
tensorrt_llm/runtime/generation.py (3 hunks)
tensorrt_llm/runtime/model_runner_cpp.py (1 hunks)
tensorrt_llm/sampling_params.py (2 hunks)
tensorrt_llm/serve/openai_protocol.py (4 hunks)

🧰 Additional context used

📓 Path-based instructions (8)

**/*.{h,hpp,hh,hxx,cpp,cxx,cc,cu,cuh,py}