[TRTLLM-13776][feat] Add Piecewise CUDA Graph support for Qwen3.5/Qwen3.6 MoE model (NVIDIA#15900)

Signed-off-by: nv-guomingz <137257613+nv-guomingz@users.noreply.github.com>

Author: nv-guomingz

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -3995,7 +3995,11 @@ def previous_seq_slots_device():
             self.previous_kv_lens_offsets_cuda *= 0
 
         position_ids = self._apply_position_id_offset(position_ids)
-        if self.use_mrope and mrope_position_ids:
+        # Use the (3,1,N) MRoPE layout whenever the model declares MRoPE, even
+        # for text-only batches: keeping position_ids rank-consistent between
+        # warmup and serving keeps torch.compile guards stable, so piecewise
+        # CUDA graphs captured at warmup remain usable at runtime.
+        if self.use_mrope:
             # Mixed batches may have only some requests with multimodal MRoPE
             # data. Seed the full (3,1,N) buffer from scalar position_ids
             # (text-only tokens get the same value on all 3 axes), then
@@ -4154,7 +4158,10 @@ def previous_seq_slots_device():
         if attn_metadata.padded_num_tokens is not None:
             self.input_ids_cuda[total_num_tokens:padded_num_tokens].fill_(0)
             virtual_num_tokens = padded_num_tokens
-            if self.use_mrope and mrope_position_ids:
+            # Match the rank of the unpadded branch: MRoPE models always use
+            # the (3,1,N) layout (see the seeding block above), so the padded
+            # view must stay 3D as well to keep torch.compile guards stable.
+            if self.use_mrope:
                 # Zero-fill padding on dim 2 (token dim) of (3,1,N) buffer.
                 self.mrope_position_ids_cuda[:, :, total_num_tokens:
                                              padded_num_tokens].fill_(0)