feat: update flashinfer ar oneshot params (#8687)
This commit is contained in:
@@ -443,7 +443,6 @@ class CommunicateWithAllReduceAndLayerNormFn:
|
||||
and _is_flashinfer_available
|
||||
and hasattr(layernorm, "forward_with_allreduce_fusion")
|
||||
and global_server_args_dict["enable_flashinfer_allreduce_fusion"]
|
||||
and hidden_states.shape[0] <= 128
|
||||
):
|
||||
hidden_states, residual = layernorm.forward_with_allreduce_fusion(
|
||||
hidden_states, residual
|
||||
|
||||
@@ -125,7 +125,7 @@ def flashinfer_allreduce_residual_rmsnorm(
|
||||
weight: torch.Tensor,
|
||||
eps: float = 1e-6,
|
||||
max_token_num: int = 128,
|
||||
use_oneshot: bool = True,
|
||||
use_oneshot: Optional[bool] = None,
|
||||
trigger_completion_at_end: bool = False,
|
||||
fp32_acc: bool = False,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
|
||||
Reference in New Issue
Block a user