Summary by CodeRabbit

• New Features
  • Improved compiler optimization for distributed multi-GPU inference. Enhanced handling of normalization operations during model compilation provides performance improvements for multi-GPU deployments.

Description

In multi-GPU mode (world_size > 1), the torch.compile custom Backend uses recover_pass to split fused_add_rmsnorm into separate add + rmsnorm ops, then re-fuses them via pattern matching. However, only allreduce-based fusion patterns (register_ar_fusions) were registered. The simpler register_add_norm fallback was missing.

This causes layers without an explicit allreduce before add + rmsnorm to remain permanently unfused — resulting in 2 kernel launches instead of 1, with ~24-30% more GPU time for these operations.

Affected scenarios:

• enable_attention_dp=True — attention layers skip allreduce (Llama, Qwen3, DeepSeekV3, etc.)
• Hybrid Mamba+Transformer models (Nemotron-H) — Mamba layers have allreduce hidden inside Linear, not as FX graph node
• Pipeline-parallel-only setups (PP > 1, TP = 1) — mpi_world_size() > 1 but no allreduce anywhere
• Fusion config flags (PRE_MLP_FUSION, POST_MLP_FUSION) — disable allreduce on specific layers

Fix: Add register_add_norm as a fallback PatternMatcherPass in the multi-GPU branch, placed after all allreduce-based fusion passes to ensure correct priority (2 lines added in backend.py).

Test Coverage

Unit tests:

• Existing tests/unittest/_torch/compilation/test_add_norm.py — 4 tests passed (single-GPU register_add_norm correctness, both dtypes, with/without inductor). Confirms no regression.

nsys profiling verification (2x NVIDIA B200):

All three testable scenarios verified with before/after nsys profiling using trtllm-serve:

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.