[Feature]: [AutoDeploy] Enable SuperV3 partial loading for fast debug

🚀 The feature, motivation and pitch

Requires applying #10319 + #10397 .

Adding this to super_v3.yaml should allow for running a partial model.

model_kwargs:
  num_hidden_layers: 4

However, this fails due to OOM.
Run:

python benchmarks/cpp/prepare_dataset.py --stdout --tokenizer /lustre/fsw/portfolios/llmservice/users/lvega/models/super-v3-iter_0440000/hf token-norm-dist --input-mean 1k --output-mean 1k --input-stdev 0 --output-stdev 0 --num-requests 16 > /tmp/super_16.inp

trtllm-bench --model NVIDIA-Super-v3-dev --model_path /lustre/fsw/portfolios/llmservice/users/lvega/models/super-v3-iter_0440000/hf throughput --backend _autodeploy --dataset /tmp/super_256.inp --tp 1 --extra_llm_api_options examples/auto_deploy/super_v3.yaml --max_batch_size 16 --max_num_tokens 2048

Specify in the command line
--max_batch_size 16 --max_num_tokens 2048
to avoid failure in calc_engine_setting which doesn't take num_hidden_layers into account.
Alternatively, need to override layers number and total params count in the else clause in tensorrt_llm/bench/benchmark/utils/general.py:L107:

    if params.get("max_batch_size") and params.get("max_num_tokens"):
        logger.info("Use user-provided max batch size and max num tokens.")
        max_batch_size, max_num_tokens = params.get(
            "max_batch_size"), params.get("max_num_tokens")
    else:
        model_config = get_model_config(model, model_path)

OOM in KV-cache allocation. Presumably, the allocation does not take into account the reduced number of layers.

[01/05/2026-03:02:48] [TRT-LLM AUTO-DEPLOY] [RANK 0] [I] [stage=cache_init, transform=initialize_cache] Initialized 4 caches for cached attention
[01/05/2026-03:02:48] [TRT-LLM AUTO-DEPLOY] [RANK 0] [I] [stage=cache_init, transform=initialize_cache] enabled=True, num_matches=4, is_clean=True, has_valid_shapes=False
[01/05/2026-03:02:48] [TRT-LLM AUTO-DEPLOY] [RANK 0] [I] [stage=cache_init, transform=initialize_cache] elapsed time: total=0.000s, pre_cleanup=0.000s, apply=0.000s, post_cleanup=0.000s
[01/05/2026-03:02:48] [TRT-LLM AUTO-DEPLOY] [RANK 0] [I] [stage=cache_init, transform=resize_kv_cache] Free memory (MB): 63165, Total memory (MB): 81007
[01/05/2026-03:02:48] [TRT-LLM AUTO-DEPLOY] [RANK 0] [I] [stage=cache_init, transform=resize_kv_cache] Current cache size (MB): 1585, Current num pages: 416
[01/05/2026-03:02:48] [TRT-LLM AUTO-DEPLOY] [RANK 0] [I] [stage=cache_init, transform=resize_kv_cache] Current KV-only cache size (MB): 0
[01/05/2026-03:02:48] [TRT-LLM AUTO-DEPLOY] [RANK 0] [I] [stage=cache_init, transform=resize_kv_cache] Free memory before forward pass (MB): 63165
[01/05/2026-03:02:48] [TRT-LLM] [I] PDL enabled
[01/05/2026-03:02:48] [TRT-LLM] [W] [AutoTuner] trtllm::fused_moe::gemm1 using the fallback tactic, due to cache miss on input shapes=(torch.Size([1920, 1024]), torch.Size([512, 2688, 1024]), torch.Size([0]), torch.Size([512, 1024, 2688]), torch.Size([0]))
[01/05/2026-03:02:48] [TRT-LLM] [W] [AutoTuner] trtllm::fused_moe::gemm2 using the fallback tactic, due to cache miss on input shapes=(torch.Size([1920, 1024]), torch.Size([512, 2688, 1024]), torch.Size([0]), torch.Size([512, 1024, 2688]), torch.Size([0]))
[01/05/2026-03:02:55] [TRT-LLM AUTO-DEPLOY] [RANK 0] [I] [stage=cache_init, transform=resize_kv_cache] Free memory after forward pass (MB): 61111
[01/05/2026-03:02:55] [TRT-LLM AUTO-DEPLOY] [RANK 0] [I] [stage=cache_init, transform=resize_kv_cache] Memory for forward pass (MB): 2054
[01/05/2026-03:02:55] [TRT-LLM AUTO-DEPLOY] [RANK 0] [I] [stage=cache_init, transform=resize_kv_cache] After all_gather - new_num_pages: 53849515622
[01/05/2026-03:02:55] [TRT-LLM AUTO-DEPLOY] [RANK 0] [I] [stage=cache_init, transform=resize_kv_cache] Resizing cache_loc capacity from 393600 to 53849515622 to accommodate num_pages=53849515622
[01/05/2026-03:02:55] [TRT-LLM] [E] Failed to initialize executor on rank 0: CUDA out of memory. Tried to allocate 200.61 GiB. GPU 0 has a total capacity of 79.11 GiB of which 61.62 GiB is free. Process 2301681 has 592.00 MiB memory in use. Including non-PyTorch memory, this process has 16.90 GiB memory in use. Of the allocated memory 15.00 GiB is allocated by PyTorch, and 698.86 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)

Alternatives

No response

Additional context

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and checked the documentation and examples for answers to frequently asked questions.