---
id: pr-tensorrt-llm-12847
repo: NVIDIA/TensorRT-LLM
pr: 12847
title: '[None][fix] Fix multi_stream_moe accuracy with MLIR and piecewise cudagraphs'
author: suyoggupta
date: '2026-04-11'
url: https://github.com/NVIDIA/TensorRT-LLM/pull/12847
source_category: upstream-code
architectures:
- sm100
tags:
- cuda kernel
- moe kernel
- triton kernel
techniques: []
hardware_features: []
kernel_types:
- moe
languages:
- cuda-cpp
- triton
captured_at: '2026-05-16'
status: merged
merge_sha: 72dc7eccecbe3ce6f542019d038e504f035ef3cd
inclusion_reason: Search-matched CUDA kernel optimization PR; materialize diff/source before synthesis.
changed_paths:
- examples/auto_deploy/model_registry/configs/gemma4_moe.yaml
- tensorrt_llm/_torch/auto_deploy/compile/backends/torch_cudagraph.py
- tensorrt_llm/_torch/auto_deploy/compile/piecewise_utils.py
- tensorrt_llm/_torch/auto_deploy/transform/library/gather_logits_before_lm_head.py
- tensorrt_llm/_torch/auto_deploy/utils/multi_stream_utils.py
- tests/unittest/auto_deploy/singlegpu/compile/test_piecewise_utils.py
- tests/unittest/auto_deploy/singlegpu/transformations/library/test_gather_logits_before_lm_head.py
artifact_dir: evidence/pull-bundles/tensorrt-llm/gh-12847
head_ref_oid: 7d0215129fe73a4c56e81aff366f81fe4fdd48f0
---
# [None][fix] Fix multi_stream_moe accuracy with MLIR and piecewise cudagraphs

- Repository: `NVIDIA/TensorRT-LLM`
- PR: [12847](https://github.com/NVIDIA/TensorRT-LLM/pull/12847)
- Decision: `Search-matched CUDA kernel optimization PR; materialize diff/source before synthesis.`
- Reason: Search-matched CUDA kernel optimization PR; materialize diff/source before synthesis.

## Changed Paths

- `examples/auto_deploy/model_registry/configs/gemma4_moe.yaml`
- `tensorrt_llm/_torch/auto_deploy/compile/backends/torch_cudagraph.py`
- `tensorrt_llm/_torch/auto_deploy/compile/piecewise_utils.py`
- `tensorrt_llm/_torch/auto_deploy/transform/library/gather_logits_before_lm_head.py`
- `tensorrt_llm/_torch/auto_deploy/utils/multi_stream_utils.py`
- `tests/unittest/auto_deploy/singlegpu/compile/test_piecewise_utils.py`
- `tests/unittest/auto_deploy/singlegpu/transformations/library/test_gather_logits_before_lm_head.py`

## Evidence Bundle

`evidence/pull-bundles/tensorrt-llm/gh-12847`
