[Bug Fix] Fix naive_block_assignment always defaulting to False due to arg misalignment (#33848)

Signed-off-by: Runkai Tao <rt572@physics.rutgers.edu>
This commit is contained in:
Runkai Tao
2026-02-11 22:30:57 -05:00
committed by GitHub
parent ec12d39d44
commit e1d97c38f8
2 changed files with 2 additions and 1 deletions

View File

@@ -219,7 +219,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
self.max_loras,
self.adapter_enabled,
expert_map,
naive_block_assignment,
naive_block_assignment=naive_block_assignment,
)
moe_state_dict["sorted_token_ids_lora"] = sorted_token_ids_lora

View File

@@ -458,6 +458,7 @@ class PunicaWrapperBase(PunicaWrapperABC):
adapter_enabled: torch.Tensor,
expert_map: torch.Tensor | None = None,
pad_sorted_ids: bool = False,
naive_block_assignment: bool = False,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
"""
Aligns tokens and experts into block-sized chunks for LoRA-based