feat: Add pp support for hybrid attn/mamba model

[yuxianq]

This PR adds pp support for the hybrid attn/mamba model by introducing layer masks to KV cache manager.

• KVCacheManager/MambaCacheManager/MambaHybridCacheManager support layer_mask now, so that each manager can manage selected layers only.
• Make MambaHybridLayer/NemotronHLayer inherits from DecoderLayer to support pp.
• Use absolute layer_idx instead of relative layer_idx in the modeling files, we will translate it to the relative one inside KV cache manager.

[yuxianq]

/bot run --disable-fail-fast --add-multi-gpu-test

[tensorrt-cicd]

PR_Github #5461 [ run ] triggered by Bot

[tensorrt-cicd]

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

[amukkara]

LGTM

[vegaluisjose]

Thanks @yuxianq , this looks great.

It would be great, however, to see if quality changes with this feature. We have another PR #4147 with this test.

@suyoggupta

[suyoggupta]

could you please add a brief PR description?

[Copilot]

Pull Request Overview

This PR adds pp support for the hybrid attn/mamba model by introducing layer masks and updating related interfaces.

• Adds generation of layer_mask and mamba_layer_mask in tests for configuration consistency.
• Updates function signatures and inheritance (e.g. NemotronHLayer now inherits from DecoderLayer and requires position_ids).
• Modifies cache manager initializations to propagate new layer_mask parameters.

Reviewed Changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
tests/unittest/_torch/modeling/test_modeling_nemotron_h.py	Adds layer_mask and mamba_layer_mask for testing hybrid override patterns.
tensorrt_llm/_torch/pyexecutor/resource_manager.py	Updates get_pp_layers and init to accept and process layer_mask.
tensorrt_llm/_torch/pyexecutor/_util.py	Adjusts cache manager creation to include the new layer_mask parameters.
tensorrt_llm/_torch/models/modeling_nemotron_h.py	Changes NemotronHLayer to inherit from DecoderLayer and updates layer indexing and forward signature with position_ids.
tensorrt_llm/_torch/models/modeling_mamba_hybrid.py	Updates MambaHybridLayer and cache manager initializations with layer_mask support.
examples/pytorch/out_of_tree_example/modeling_opt.py	Updates forward methods to include position_ids for decoder layers.

Comments suppressed due to low confidence (1)

tensorrt_llm/_torch/models/modeling_nemotron_h.py:173

• The new simple enumeration for layer indexing replaces previously separate counters for each layer type; please verify that this change in indexing logic is intentional and that it correctly supports layer-specific configurations.

for layer_idx, layer_type in enumerate(config.hybrid_override_pattern):

[suyoggupta]

Added copilot to also review this PR

[yuxianq]

@vegaluisjose I have cherry-pick #4147 to this PR and the new test can pass locally.

[yuxianq]

could you please add a brief PR description?

@suyoggupta Added.

[yuxianq]

/bot reuse-pipeline

[tensorrt-cicd]

PR_Github #5681 [ reuse-pipeline ] triggered by Bot

[tensorrt-cicd]

PR_Github #5681 [ reuse-pipeline ] completed with state SUCCESS
Reusing PR_Github #5461 for commit ba33051