[Flashinfer] Integrate flashinfer router gemm for sm103 (#23285)
This commit is contained in:
@@ -335,7 +335,7 @@ class MoEGate(nn.Module):
|
|||||||
and (self.weight.shape[0] == 256 or self.weight.shape[0] == 384)
|
and (self.weight.shape[0] == 256 or self.weight.shape[0] == 384)
|
||||||
and _device_sm >= 90
|
and _device_sm >= 90
|
||||||
):
|
):
|
||||||
if _device_sm == 100 and self.weight.shape[0] == 256:
|
if _device_sm in [100, 103] and self.weight.shape[0] == 256:
|
||||||
# router gemm output float32
|
# router gemm output float32
|
||||||
logits = torch.empty(
|
logits = torch.empty(
|
||||||
hidden_states.shape[0],
|
hidden_states.shape[0],
|
||||||
|
|||||||
Reference in New Issue
Block a user