[TRTLLM-5252][fix] Propagate mapping to intermediate layers

Signed-off-by: William Zhang <[email protected]>

Author: 2ez4bz

examples/llm-api/quickstart_advanced.py

@@ -150,7 +150,9 @@ def parse_arguments():
 def setup_llm(args, **kwargs):
     kv_cache_config = KvCacheConfig(
         enable_block_reuse=not args.disable_kv_cache_reuse,
-        free_gpu_memory_fraction=args.kv_cache_fraction,
+        # free_gpu_memory_fraction=args.kv_cache_fraction,
+        free_gpu_memory_fraction=0.5,
+        max_tokens=10_000,
         dtype=args.kv_cache_dtype,
     )
 

tensorrt_llm/_torch/models/modeling_mistral.py

@@ -475,6 +475,7 @@ def __init__(self, model_config: ModelConfig[Mistral3Config]):
             out_features=hidden_size,
             bias=False,
             dtype=config.torch_dtype,
+            mapping=model_config.mapping,
         )
 
     @torch.inference_mode()
@@ -546,6 +547,7 @@ def __init__(self, model_config: ModelConfig[Mistral3Config]):
             out_features=config.text_config.hidden_size,
             bias=config.multimodal_projector_bias,
             dtype=dtype,
+            mapping=model_config.mapping,
         )
 
     @torch.inference_mode()

tests/integration/test_lists/test-db/l0_dgx_h100.yml

@@ -51,6 +51,7 @@ l0_dgx_h100:
   - accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_ctx_pp_gen_tp_asymmetric[MMLU-gen_tp=1-ctx_pp=2]
   - accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_ctx_pp_gen_tp_asymmetric[MMLU-gen_tp=2-ctx_pp=2]
   - test_e2e.py::test_ptp_quickstart_advanced_bs1
+  - unittest/_torch/modeling/test_modeling_pixtral.py::test_tensor_parallelism
 - condition:
     ranges:
       system_gpu_count: