Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
pggPL wants to merge 1 commit into
base: main
Choose a base branch
from
+32 −10
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 19 additions & 0 deletions tests/pytorch/test_permutation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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])
Expand Down
23 changes: 13 additions & 10 deletions transformer_engine/pytorch/permutation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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}.")
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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}.")
Expand Down Expand Up @@ -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
Expand Down
Loading