[PyTorch] Fix fake kernel metadata mismatch for empty-input MoE permute ops (#3444)

Fix fake kernel metadata mismatch for empty-input MoE permute ops

For empty input, moe_permute (mask and index map) returned row_id_map
created with torch.tensor([]) (float32), while the registered fake
kernels declare int32 with different shapes. PyTorch >= 2.14 inductor
validates fallback output dtype metadata at runtime (pytorch#183731),
so torch.compile'd empty-input permute now errors with
'expected dtype torch.int32 but got torch.float32'.

Return correctly-typed empty tensors from the eager fast paths, align
the index-map fake output shape with the empty-input fast path, and
keep permuted_probs dtype consistent between eager and fake.
Add an empty-input test for the index map path.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Author: pggPL

tests/pytorch/test_permutation.py

@@ -1705,6 +1705,25 @@ def test_permutation_index_map(
     )
 
 
+@pytest.mark.parametrize("te_dtype", _te_dtypes)
+@pytest.mark.parametrize("use_torch_compile", [False, True])
+def test_permutation_index_map_empty_input(te_dtype, use_torch_compile):
+    with_probs = True
+    BENCHMARK = False
+
+    _test_permutation_index_map(
+        te_dtype=te_dtype,
+        num_tokens=0,
+        num_expert=8,
+        hidden_size=4096,
+        topK=2,
+        num_out_tokens=0,
+        with_probs=with_probs,
+        BENCHMARK=BENCHMARK,
+        use_torch_compile=use_torch_compile,
+    )
+
+
 @pytest.mark.parametrize("te_dtype", _te_dtypes)
 @pytest.mark.parametrize("num_tokens", [4096])
 @pytest.mark.parametrize("num_expert", [7, 16])

transformer_engine/pytorch/permutation.py

@@ -42,7 +42,7 @@ def moe_permute_index_map_forward(
     global _moe_permute_index_map_workspace, _moe_permute_index_map_max_expanded_token_num
 
     if not inp.numel():
-        return inp.clone(), torch.tensor([], device=inp.device)
+        return inp.clone(), torch.empty(0, dtype=torch.int32, device=inp.device)
 
     if not inp.is_cuda:
         raise ValueError(f"inp must be a CUDA tensor, but got tensor on {inp.device}.")
@@ -98,9 +98,10 @@ def _moe_permute_index_map_fake(  # pylint: disable=unused-argument
         assert (
             num_out_tokens >= 0
         ), f"moe_permute (index map) requires num_out_tokens >= 0, got {num_out_tokens}."
-
-    # Infer output shape
-    output_tokens = num_out_tokens if num_out_tokens > 0 else num_tokens * topK
+        output_tokens = num_out_tokens if num_out_tokens > 0 else num_tokens * topK
+    else:
+        # Match `moe_permute_index_map_forward` empty-input fast path (ignores num_out_tokens).
+        output_tokens = 0
 
     # row_id_map is 1D with size = num_tokens * topK
     fake_output = torch.empty((output_tokens, inp.shape[1]), dtype=inp.dtype, device=inp.device)
@@ -290,7 +291,13 @@ def moe_permute_mask_map_forward(
 ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     """Forward pass for MoE permute with mask router map."""
     if not inp.numel():
-        return inp.clone(), torch.tensor([], device=inp.device), torch.tensor([], device=inp.device)
+        row_id_map = torch.empty(
+            (0, routing_map.size(1) * 2 + 1), dtype=torch.int32, device=inp.device
+        )
+        permuted_probs = torch.empty(
+            0, dtype=probs.dtype if probs is not None else torch.float32, device=inp.device
+        )
+        return inp.clone(), row_id_map, permuted_probs
 
     if not inp.is_cuda:
         raise ValueError(f"inp must be a CUDA tensor, but got tensor on {inp.device}.")
@@ -450,11 +457,7 @@ def _moe_permute_mask_map_forward_fake(  # pylint: disable=unused-argument
         (num_tokens, num_experts * 2 + 1), dtype=torch.int32, device=inp.device
     )
     if probs is not None:
-        fake_permuted_probs = (
-            torch.empty((out_rows,), dtype=probs.dtype, device=inp.device)
-            if out_rows > 0
-            else torch.empty(0, device=inp.device)
-        )
+        fake_permuted_probs = torch.empty((out_rows,), dtype=probs.dtype, device=inp.device)
     else:
         fake_permuted_probs = torch.empty(0, device=inp.device)
     return fake_output, fake_row_id_map, fake_permuted_probs