[nvbug/5280806][fix] Fix 2 model spec decode flow

[mikeiovine]

[fix] Fix 2 model spec decode flow

Description

Fix a small issue introduced by the feat/llama4 integration (#4739).

There was a bad merge:

• In [fix] Fix a few issues with EAGLE3 in PyTorch backend #3686, we did some code cleanup on the 2 engine path. This line was confusing because of the hardcoded layer index:

spec_metadata.maybe_capture_hidden_states(1, hidden_states_to_save)

The fix was to capture the hidden states in the midlayer, which makes everything consistent with how hidden states are captured in the target model.

• The feat/llama4 branch added these lines back by mistake. They don't work any more due to an API change.

The one model flow doesn't use this code path, so it was probably untested.

Also fixes another small issue introduced by #4379. The prompt_length was not being set correctly for extend_ctx cases.

Test Coverage

Re-enabled EAGLE3 2 model tests in L0.

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 [--disable-fail-fast --skip-test --stage-list "A10-1, xxx" --gpu-type "A30, H100_PCIe" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-[Post-Merge]-1, xxx"]

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

--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-1, xxx" (OPTIONAL) : Only run the specified test stages. Examples: "A10-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.

--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. Will also run 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-[Post-Merge]-1, xxx" (OPTIONAL) : Run the ordinary L0 pre-merge pipeline and specified test stages. Examples: --extra-stage "H100_PCIe-[Post-Merge]-1, xxx".

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.

[mikeiovine]

/bot run

[tensorrt-cicd]

PR_Github #7215 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #7215 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #5225 completed with status: 'FAILURE'

[mikeiovine]

/bot run

[tensorrt-cicd]

PR_Github #7233 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #7233 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #5238 completed with status: 'FAILURE'

[mikeiovine]

Fixed one more issue (extend_ctx method using isinstance instead of issubclass, causing the prefill kernels to be used instead of the new generation kernels on Blackwell).

Also uncovered one more bug. Added a workaround in the unit tests and opened https://nvbugspro.nvidia.com/bug/5314469 to follow up.

[mikeiovine]

/bot run

[tensorrt-cicd]

PR_Github #7240 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #7240 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #5244 completed with status: 'FAILURE'

[mikeiovine]

/bot run

[tensorrt-cicd]

PR_Github #7357 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #7357 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #5333 completed with status: 'FAILURE'

[mikeiovine]

/bot run

[tensorrt-cicd]

PR_Github #7376 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #7376 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #5351 completed with status: 'FAILURE'

[mikeiovine]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #7530 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #7530 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #5462 completed with status: 'FAILURE'

[mikeiovine]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #7742 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #7742 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #5611 completed with status: 'FAILURE'

[mikeiovine]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #7792 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #7792 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #5643 completed with status: 'FAILURE'

[mikeiovine]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #7915 [ run ] triggered by Bot

[mikeiovine]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #7936 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #7915 [ run ] completed with state ABORTED

[tensorrt-cicd]

PR_Github #7936 [ run ] completed with state FAILURE
/LLM/main/L0_MergeRequest_PR pipeline #5739 completed with status: 'FAILURE'

[mikeiovine]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #7990 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #7990 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #5789 completed with status: 'SUCCESS'