perf: accelerate decode linear with bf16 gemv
- add decode-shape benchmark harness - add bf16 GEMV CUDA primitive with head-dim generic kernel - dispatch decode-time linear layers to gemv for M=1 - extend gemv coverage to small decode batches
This commit is contained in:
@@ -1,8 +1,9 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from torch import Tensor
|
||||
|
||||
from astrai.extension.backend.linear import linear
|
||||
|
||||
|
||||
class Linear(nn.Module):
|
||||
def __init__(
|
||||
@@ -21,4 +22,4 @@ class Linear(nn.Module):
|
||||
nn.init.uniform_(self.bias, -bound, bound)
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
return F.linear(x, self.weight, self.bias)
|
||||
return linear(x, self.weight, self.bias)
|
||||
|
||||
Reference in New Issue
Block a user