[llama4] add back max_len arg to generate_permute_indices (#1354)

danielvegamyhre · web-flow · commit 05d3f7cb751d · 2025-06-27T10:21:17.000-07:00
Fixes #1353 Bug is related to this revert, the arg list changed: #1340 So need to add this arg back: 4e31b86 cc @lessw2020 @tianyu-l
diff --git a/torchtitan/experiments/llama4/model/moe.py b/torchtitan/experiments/llama4/model/moe.py
@@ -295,6 +295,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
                     num_local_tokens_per_expert,
                     self.experts.num_experts,
                     1,
+                    token_indices.shape[0] + self.experts.num_experts * ALIGN_SIZE_M,
                     ALIGN_SIZE_M,
                 )
             token_indices = torch.vstack(

Original file line number	Diff line number	Diff line change
`@@ -295,6 +295,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:`
`295`	`295`	`num_local_tokens_per_expert,`
`296`	`296`	`self.experts.num_experts,`
`297`	`297`	`1,`
	`298`	`+ token_indices.shape[0] + self.experts.num_experts * ALIGN_SIZE_M,`
`298`	`299`	`ALIGN_SIZE_M,`
`299`	`300`	`)`
`300`	`301`	`token_indices = torch.vstack(`