EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

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

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files

7
python/sglang/srt/layers/activation.py
View File

@@ -30,18 +30,11 @@ from vllm.model_executor.utils import set_weight_attrs
class SiluAndMul(CustomOp):
def __init__(self, **kwargs):
super().__init__()
self.is_lower_sm80 = torch.cuda.get_device_capability()[0] < 8
def forward_native(self, x: torch.Tensor) -> torch.Tensor:
d = x.shape[-1] // 2
return F.silu(x[..., :d]) * x[..., d:]
def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
if self.is_lower_sm80:
return self.forward_native(x)
d = x.shape[-1] // 2
output_shape = x.shape[:-1] + (d,)
out = torch.empty(output_shape, dtype=x.dtype, device=x.device)