[AutoDeploy] Nemo SuperV3: Enable multistream MoE

[galagam]

Nemotron Super V3 crashes when multisteam MoE transformation is enabled.
Enable and check impact on perf.

  multi_stream_moe:
    stage: compile
    enabled: true # 251230 fails

Trace:

[12/30/2025-06:30:13] [TRT-LLM AUTO-DEPLOY] [RANK 1] [I] [stage=compile, transform=fuse_causal_conv_activation] enabled=True, num_matches=40, is_clean=True, has_valid_shapes=False
[12/30/2025-06:30:13] [TRT-LLM AUTO-DEPLOY] [RANK 1] [I] [stage=compile, transform=fuse_causal_conv_activation] elapsed time: total=0.141s, pre_cleanup=0.000s, apply=0.070s, post_cleanup=0.071s
[12/30/2025-06:30:13] [TRT-LLM] [RANK 1] [E] Failed to initialize executor on rank 1: Target input node not found for node trtllm_moe_fused
[12/30/2025-06:30:13] [TRT-LLM] [RANK 1] [E] Traceback (most recent call last):
  File "/home/scratch.ghubaraagam_sw_2/TensorRT-LLM_github/tensorrt_llm/executor/worker.py", line 264, in worker_main
    worker: GenerationExecutorWorker = worker_cls(
                                       ^^^^^^^^^^^
  File "/home/scratch.ghubaraagam_sw_2/TensorRT-LLM_github/tensorrt_llm/executor/worker.py", line 60, in __init__
    self.setup_engine()
  File "/home/scratch.ghubaraagam_sw_2/TensorRT-LLM_github/tensorrt_llm/executor/base_worker.py", line 257, in setup_engine
    self.engine = _create_py_executor(
                  ^^^^^^^^^^^^^^^^^^^^
  File "/home/scratch.ghubaraagam_sw_2/TensorRT-LLM_github/tensorrt_llm/executor/base_worker.py", line 228, in _create_py_executor
    _executor = create_executor(**args)
                ^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/scratch.ghubaraagam_sw_2/TensorRT-LLM_github/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py", line 1026, in create_autodeploy_executor
    engine = ADEngine.build_from_config(ad_config=ad_config, mapping=dist_mapping)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/scratch.ghubaraagam_sw_2/TensorRT-LLM_github/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py", line 449, in build_from_config
    return cls(
           ^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/utils/_contextlib.py", line 120, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/scratch.ghubaraagam_sw_2/TensorRT-LLM_github/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py", line 519, in __init__
    self.model = get_inference_model(self.cache_seq_interface)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/scratch.ghubaraagam_sw_2/TensorRT-LLM_github/tensorrt_llm/_torch/auto_deploy/transform/optimizer.py", line 71, in __call__
    mod = transform(mod, cm, self.factory, self.shared_config)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/scratch.ghubaraagam_sw_2/TensorRT-LLM_github/tensorrt_llm/_torch/auto_deploy/transform/interface.py", line 213, in _wrapper
    return call_fn(self, gm, cm, factory, shared_config)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/scratch.ghubaraagam_sw_2/TensorRT-LLM_github/tensorrt_llm/_torch/auto_deploy/transform/interface.py", line 359, in __call__
    mod, info_apply = self._apply_per_gm_or_whole_model(mod, cm, factory, shared_config)
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/scratch.ghubaraagam_sw_2/TensorRT-LLM_github/tensorrt_llm/_torch/auto_deploy/transform/interface.py", line 418, in _apply_per_gm_or_whole_model
    graph_sub, info_apply = self._apply(graph_sub, cm, factory, shared_config)
                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/scratch.ghubaraagam_sw_2/TensorRT-LLM_github/tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_moe.py", line 312, in _apply
    gm, num_matches = _execute_op_in_aux_stream(gm, op_dict)
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/scratch.ghubaraagam_sw_2/TensorRT-LLM_github/tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_moe.py", line 270, in _execute_op_in_aux_stream
    assert target_input_node is not None, f"Target input node not found for node {n}"
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AssertionError: Target input node not found for node trtllm_moe_fused

[galagam]

Config, copied from nanov3 with necessary changes:

world_size: 4
compile_backend: torch-cudagraph
max_batch_size: 512
max_seq_len: 2048
skip_loading_weights: true
free_mem_ratio: 0.8
cuda_graph_batch_sizes: [1, 2, 4, 8, 16, 32, 64, 128, 256, 512]
kv_cache_config:
  # disable kv_cache reuse since not supported for hybrid/ssm models
  enable_block_reuse: false
transforms:
  detect_sharding:
    allreduce_strategy: SYMM_MEM
    sharding_dims: ['ep', 'bmm']
    manual_config:
      head_dim: 128
      tp_plan:
        # mamba SSM layer
        "in_proj": "mamba"
        "out_proj": "rowwise"
        # attention layer
        "q_proj": "colwise"
        "k_proj": "colwise"
        "v_proj": "colwise"
        "o_proj": "rowwise"
        # NOTE: consider not sharding shared experts and/or
        # latent projections at all, keeping them replicated.
        # To do so, comment out the corresponding entries.
        # moe layer: SHARED experts
        "up_proj": "colwise"
        "down_proj": "rowwise"
        # MoLE: latent projections: simple shard
        "fc1_latent_proj": "gather"
        "fc2_latent_proj": "gather"
  multi_stream_moe:
    stage: compile
    enabled: false # 251230 fails
  # tunable mamba cache dtype
  # --> use float32 for accuracy and default (null) for speed
  insert_cached_ssm_attention:
      cache_config:
        # mamba_dtype: float32
        mamba_dtype: null
  gather_logits_before_lm_head:
    # TODO: fix https://github.com/NVIDIA/TensorRT-LLM/issues/9878 to enable by default
    enabled: true
  fuse_mamba_a_log:
    stage: post_load_fusion
    enabled: true

[galagam]

On ToT:
Edit examples/auto_deploy/super_v3.yaml to enable multi_stream_moe, then run:

trtllm-bench --model NVIDIA-Super-v3-dev --model_path /home/scratch.ghubaraagam_sw_2/models/super-v3-iter_0440000/hf throughput --backend _autodeploy --dataset debug/super/synthetic_1k1k_256req.txt --tp 4 --extra_llm_api_options examples/auto_deploy/super_v3.yaml