[None][fix] SpecDecOneEngineForCausalLM: accept optional hidden_size/vocab_size for composite configs

[brnguyen2]

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__ to accept optional hidden_size and vocab_size parameters.
• When explicit values are provided, they are forwarded to DecoderModelForCausalLM.__init__; otherwise, the constructor resolves hidden_size/vocab_size from model_config.pretrained_config (preserving existing behavior for current callers).
• No configuration files or tests/integration/test_lists/ entries were modified.

QA Engineer Review

• Added 3 unit tests in tests/unittest/_torch/modeling/test_modeling_speculative.py for SpecDecOneEngineForCausalLM.__init__:
  • test_specdec_one_engine_reads_from_pretrained_config (default lookup from pretrained_config)
  • test_specdec_one_engine_accepts_explicit_sizes (explicit values when pretrained_config lacks top-level attributes)
  • test_specdec_one_engine_explicit_overrides_pretrained_config (explicit overrides take precedence)
• The tests mock DecoderModelForCausalLM.__init__ and assert the expected forwarded hidden_size/vocab_size.
• No integration test-list entries in tests/integration/test_lists/ were changed/added for this coverage.
• Verdict: insufficient.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 4704ac9b-89b2-4344-8a43-9ff1f4020a9b

📥 Commits

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

📒 Files selected for processing (2)

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

🚧 Files skipped from review as they are similar to previous changes (2)

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

---

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 change and follows the required [None][fix] format.
Description check	✅ Passed	It includes the required Description and Test Coverage sections and is sufficiently complete despite an unfilled PR Checklist.
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