[None][fix] SpecDecOneEngineForCausalLM: accept optional hidden_size/vocab_size for composite configs by brnguyen2 · Pull Request #16762 · NVIDIA/TensorRT-LLM

brnguyen2 · 2026-07-22T23:09:17Z

Description

SpecDecOneEngineForCausalLM.__init__ previously read hidden_size and vocab_size unconditionally from model_config.pretrained_config, which breaks for composite model configs (e.g. vision-language wrappers) where these fields live in a nested text_config rather than at the top level.

Two new optional constructor parameters (hidden_size, vocab_size) let callers pass the text-config values explicitly. When omitted, the existing behaviour is preserved — values are read from pretrained_config as before — so all current callers are unaffected.

Root cause

# Before: always reads from top-level pretrained_config
super().__init__(model,
                 config=model_config,
                 hidden_size=model_config.pretrained_config.hidden_size,
                 vocab_size=model_config.pretrained_config.vocab_size)

For models whose config wraps a text_config (e.g. Qwen3NextForCausalLM, which subclasses SpecDecOneEngineForCausalLM), the top-level config does not expose these attributes, causing AttributeError on construction.

Fix

Make both parameters optional with None defaults; fall through to pretrained_config only when not provided explicitly:

def __init__(self,
             model: TModel,
             model_config: ModelConfig[TConfig],
             hidden_size: Optional[int] = None,
             vocab_size: Optional[int] = None):
    if hidden_size is None:
        hidden_size = model_config.pretrained_config.hidden_size
    if vocab_size is None:
        vocab_size = model_config.pretrained_config.vocab_size
    super().__init__(model, config=model_config,
                     hidden_size=hidden_size, vocab_size=vocab_size)

Test Coverage

Added three mock-based unit tests to tests/unittest/_torch/modeling/test_modeling_speculative.py (no GPU required, <1s):

test_specdec_one_engine_reads_from_pretrained_config: default path passes values from pretrained_config
test_specdec_one_engine_accepts_explicit_sizes: composite-config path uses caller-supplied sizes when pretrained_config lacks the attributes
test_specdec_one_engine_explicit_overrides_pretrained_config: explicit args take precedence even when pretrained_config has values

Dev Engineer Review

Updated SpecDecOneEngineForCausalLM.__init__ (tensorrt_llm/_torch/models/modeling_speculative.py) to accept optional hidden_size: int | None and vocab_size: int | None.
Resolution logic: when hidden_size/vocab_size are None, they fall back to model_config.pretrained_config.hidden_size / .vocab_size; provided values are forwarded to DecoderModelForCausalLM.__init__ via super(...) (preserving prior behavior while enabling overrides for composite configs).
Added unit tests validating:
- default fallback to pretrained_config,
- forwarding of explicit constructor sizes when pretrained_config omits them,
- explicit overrides taking precedence when both are present.
Test mocking approach is consistent with the base class PostInitCaller behavior by stubbing DecoderModelForCausalLM.__init__, __post_init__, and __pp_init__.

QA Engineer Review

Touched: tests/unittest/_torch/modeling/test_modeling_speculative.py
Added test helper / functions:
- _init_specdec_with_mocked_base(model_config, **kwargs)
- test_specdec_one_engine_reads_from_pretrained_config
- test_specdec_one_engine_accepts_explicit_sizes
- test_specdec_one_engine_explicit_overrides_pretrained_config
Integration test-list coverage:
- Existing CI filtering includes unittest/_torch/modeling -k "modeling_speculative" (e.g., tests/integration/test_lists/test-db/l0_a30.yml), which should exercise the updated test_modeling_speculative.py unit tests.
- No changes were made to tests/integration/test_lists/ / tests/qa/ / tests/test-db/.
Verdict: sufficient.

coderabbitai · 2026-07-22T23:12:57Z

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

@coderabbitai resume to resume automatic reviews.
@coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

▶️ Resume reviews
🔍 Trigger review

Walkthrough

SpecDecOneEngineForCausalLM accepts optional hidden_size and vocab_size overrides, falls back to pretrained configuration values, and forwards the resolved values to its superclass. Unit tests cover fallback behavior and explicit-argument precedence.

Changes

Speculative model dimension overrides

Layer / File(s)	Summary
Constructor overrides and validation `tensorrt_llm/_torch/models/modeling_speculative.py`, `tests/unittest/_torch/modeling/test_modeling_speculative.py`	The constructor resolves dimensions from explicit arguments or pretrained configuration, while tests verify fallback behavior and explicit-argument precedence.

Estimated code review effort: 2 (Simple) | ~10 minutes

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title matches the main change and follows the required ticket/type format.
Description check	✅ Passed	The description clearly explains the issue, fix, and test coverage, and is mostly complete for the template.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

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

coderabbitai

🧹 Nitpick comments (1)

tensorrt_llm/_torch/models/modeling_speculative.py (1)
1897-1901: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add required annotations to the changed functions.

tensorrt_llm/_torch/models/modeling_speculative.py#L1897-L1901: use int | None for the new parameters and add -> None.

tests/unittest/_torch/modeling/test_modeling_speculative.py#L157-L157: annotate helper arguments and its MagicMock return type.

tests/unittest/_torch/modeling/test_modeling_speculative.py#L170-L170: add -> None.

tests/unittest/_torch/modeling/test_modeling_speculative.py#L185-L185: add -> None.

tests/unittest/_torch/modeling/test_modeling_speculative.py#L205-L205: add -> None.
Suggested change
-                 hidden_size: Optional[int] = None,
-                 vocab_size: Optional[int] = None):
+                 hidden_size: int | None = None,
+                 vocab_size: int | None = None) -> None:
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/_torch/models/modeling_speculative.py` around lines 1897 - 1901,
Annotate the changed __init__ in
tensorrt_llm/_torch/models/modeling_speculative.py:1897-1901 with int | None for
hidden_size and vocab_size and -> None. In
tests/unittest/_torch/modeling/test_modeling_speculative.py:157, annotate the
helper arguments and MagicMock return type; add -> None to the helpers at lines
170, 185, and 205.
Source: Coding guidelines

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@tensorrt_llm/_torch/models/modeling_speculative.py`:
- Around line 1897-1901: Annotate the changed __init__ in
tensorrt_llm/_torch/models/modeling_speculative.py:1897-1901 with int | None for
hidden_size and vocab_size and -> None. In
tests/unittest/_torch/modeling/test_modeling_speculative.py:157, annotate the
helper arguments and MagicMock return type; add -> None to the helpers at lines
170, 185, and 205.

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 13000814-63f8-498c-b643-6e438b9049ba

📥 Commits

Reviewing files that changed from the base of the PR and between b294868 and b4543fa.

📒 Files selected for processing (2)

tensorrt_llm/_torch/models/modeling_speculative.py
tests/unittest/_torch/modeling/test_modeling_speculative.py

brnguyen2 · 2026-07-24T14:36:09Z

/bot run

tensorrt-cicd · 2026-07-24T14:41:51Z

PR_Github #61586 [ run ] triggered by Bot. Commit: 4ffd9ed Link to invocation

tensorrt-cicd · 2026-07-24T16:12:09Z

PR_Github #61586 [ run ] completed with state FAILURE. Commit: 4ffd9ed
/LLM/main/L0_MergeRequest_PR pipeline #49797 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

brnguyen2 · 2026-07-27T18:10:45Z

/bot run

tensorrt-cicd · 2026-07-27T18:16:09Z

PR_Github #61971 [ run ] triggered by Bot. Commit: ade0e53 Link to invocation

tensorrt-cicd · 2026-07-27T21:34:54Z

PR_Github #61971 [ run ] completed with state SUCCESS. Commit: ade0e53
/LLM/main/L0_MergeRequest_PR pipeline #50158 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

brnguyen2 · 2026-07-27T21:46:12Z

/bot run

tensorrt-cicd · 2026-07-27T21:53:40Z

PR_Github #62002 [ run ] triggered by Bot. Commit: ade0e53 Link to invocation

tensorrt-cicd · 2026-07-27T23:31:57Z

PR_Github #62002 [ run ] completed with state SUCCESS. Commit: ade0e53
/LLM/main/L0_MergeRequest_PR pipeline #50189 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

brnguyen2 · 2026-07-27T23:39:35Z

/bot run

tensorrt-cicd · 2026-07-27T23:44:51Z

PR_Github #62027 [ run ] triggered by Bot. Commit: ade0e53 Link to invocation

tensorrt-cicd · 2026-07-28T00:22:19Z

PR_Github #62027 [ run ] completed with state SUCCESS. Commit: ade0e53
/LLM/main/L0_MergeRequest_PR pipeline #50214 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

…vocab_size for composite configs Composite model configs (e.g. vision-language wrappers) may store hidden_size and vocab_size inside a nested text_config rather than at the top level of pretrained_config. The two new optional parameters let callers pass the text-config values explicitly; when omitted the existing behaviour (read from pretrained_config) is preserved. Regression test: three mock-based cases covering the default path, the explicit-override path, and the no-top-level-attrs path. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>

…on, annotations - Tests use real ModelConfig/PretrainedConfig instead of MagicMock configs - Construct SpecDecOneEngineForCausalLM directly instead of __new__/__init__ - Fold return_value into patch(); expected values held in variables - Annotate the new constructor params (int | None, -> None) and tests Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>

Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>

DecoderModelForCausalLM is built on the PostInitCaller metaclass, which calls __post_init__/__pp_init__ right after __init__ returns. The new SpecDecOneEngineForCausalLM tests mocked only __init__, so the hooks ran on a half-initialized object and raised AttributeError: 'SpecDecOneEngineForCausalLM' object has no attribute 'model_config'. Stub the metaclass hooks as well. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>

brnguyen2 · 2026-07-28T00:41:43Z

/bot run

tensorrt-cicd · 2026-07-28T00:47:20Z

PR_Github #62042 [ run ] triggered by Bot. Commit: 84e392e Link to invocation

tensorrt-cicd · 2026-07-28T02:17:14Z

PR_Github #62042 [ run ] completed with state SUCCESS. Commit: 84e392e
/LLM/main/L0_MergeRequest_PR pipeline #50226 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

brnguyen2 · 2026-07-30T03:58:06Z

/bot run

tensorrt-cicd · 2026-07-30T04:04:37Z

PR_Github #62690 [ run ] triggered by Bot. Commit: b79424a Link to invocation

tensorrt-cicd · 2026-07-30T06:38:39Z

PR_Github #62690 [ run ] completed with state SUCCESS. Commit: b79424a
/LLM/main/L0_MergeRequest_PR pipeline #50827 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

brnguyen2 · 2026-07-30T13:24:33Z

/bot run

brnguyen2 · 2026-07-30T13:25:33Z

/bot run

tensorrt-cicd · 2026-07-30T13:30:47Z

PR_Github #62781 [ run ] triggered by Bot. Commit: b79424a Link to invocation

tensorrt-cicd · 2026-07-30T13:31:12Z

PR_Github #62782 [ run ] triggered by Bot. Commit: b79424a Link to invocation

tensorrt-cicd · 2026-07-30T13:31:15Z

PR_Github #62781 [ run ] completed with state ABORTED. Commit: b79424a

Link to invocation

tensorrt-cicd · 2026-07-30T16:16:12Z

PR_Github #62782 [ run ] completed with state SUCCESS. Commit: b79424a
/LLM/main/L0_MergeRequest_PR pipeline #50911 completed with status: 'SUCCESS'

CI Report

Link to invocation

brnguyen2 requested a review from a team as a code owner July 22, 2026 23:09

brnguyen2 requested a review from 2ez4bz July 22, 2026 23:09

github-actions Bot assigned brnguyen2 Jul 22, 2026

coderabbitai Bot reviewed Jul 22, 2026

View reviewed changes

2ez4bz reviewed Jul 22, 2026

View reviewed changes

2ez4bz approved these changes Jul 24, 2026

View reviewed changes

brnguyen2 added 4 commits July 27, 2026 19:41

[None][fix] Apply ruff/ruff-format fixes to test_modeling_speculative.py

fcf8e9b

Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>

brnguyen2 force-pushed the port/specdec-one-engine-composite-config branch from ade0e53 to 84e392e Compare July 28, 2026 00:41

Merge branch 'main' into port/specdec-one-engine-composite-config

b79424a

brnguyen2 merged commit 353a4ee into NVIDIA:main Jul 30, 2026
7 checks passed

brnguyen2 mentioned this pull request Jul 31, 2026

[TRTLLM-14810][chore] Catch-up merge: main into feat/kimi_k3 #17143

Closed

2 tasks

Uh oh!

Conversation

brnguyen2 commented Jul 22, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Root cause

Fix

Test Coverage

Uh oh!

coderabbitai Bot commented Jul 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Reviews paused

Walkthrough

Changes

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

brnguyen2 commented Jul 24, 2026

Uh oh!

tensorrt-cicd commented Jul 24, 2026

Uh oh!

tensorrt-cicd commented Jul 24, 2026

Uh oh!

brnguyen2 commented Jul 27, 2026

Uh oh!

tensorrt-cicd commented Jul 27, 2026

Uh oh!

tensorrt-cicd commented Jul 27, 2026

Uh oh!

brnguyen2 commented Jul 27, 2026

Uh oh!

tensorrt-cicd commented Jul 27, 2026

Uh oh!

tensorrt-cicd commented Jul 27, 2026

Uh oh!

brnguyen2 commented Jul 27, 2026

Uh oh!

tensorrt-cicd commented Jul 27, 2026

Uh oh!

tensorrt-cicd commented Jul 28, 2026

Uh oh!

brnguyen2 commented Jul 28, 2026

Uh oh!

tensorrt-cicd commented Jul 28, 2026

Uh oh!

tensorrt-cicd commented Jul 28, 2026

Uh oh!

brnguyen2 commented Jul 30, 2026

Uh oh!

tensorrt-cicd commented Jul 30, 2026

Uh oh!

tensorrt-cicd commented Jul 30, 2026

Uh oh!

brnguyen2 commented Jul 30, 2026

Uh oh!

brnguyen2 commented Jul 30, 2026

Uh oh!

tensorrt-cicd commented Jul 30, 2026

Uh oh!

tensorrt-cicd commented Jul 30, 2026

Uh oh!

tensorrt-cicd commented Jul 30, 2026

Uh oh!

tensorrt-cicd commented Jul 30, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

brnguyen2 commented Jul 22, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Jul 22, 2026 •

edited

Loading