feat: support sm75 with FlashInfer v0.1.6 (#1233)

This commit is contained in:
Yineng Zhang
2024-08-28 18:39:12 +10:00
committed by GitHub
parent 6cc38b2bf3
commit 198974cd1a
5 changed files with 4 additions and 12 deletions

View File

@@ -32,15 +32,12 @@ class RMSNorm(CustomOp):
super().__init__()
self.weight = nn.Parameter(torch.ones(hidden_size))
self.variance_epsilon = eps
self.is_lower_sm80 = torch.cuda.get_device_capability()[0] < 8
def forward_cuda(
self,
x: torch.Tensor,
residual: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
if self.is_lower_sm80:
return self.forward_native(x, residual)
if residual is not None:
fused_add_rmsnorm(x, residual, self.weight.data, self.variance_epsilon)