0z5a
a144d7f306
- 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
26 lines
843 B
Python
26 lines
843 B
Python
import torch
|
|
import torch.nn as nn
|
|
from torch import Tensor
|
|
|
|
from astrai.extension.backend.linear import linear
|
|
|
|
|
|
class Linear(nn.Module):
|
|
def __init__(
|
|
self, in_dim: int, out_dim: int, bias: bool = False, init_std: float = 0.02
|
|
):
|
|
super().__init__()
|
|
self.weight = nn.Parameter(torch.empty((out_dim, in_dim)))
|
|
self.bias = nn.Parameter(torch.zeros(out_dim)) if bias else None
|
|
self.init_std = init_std
|
|
|
|
def reset_parameters(self):
|
|
nn.init.normal_(self.weight, mean=0.0, std=self.init_std)
|
|
if self.bias is not None:
|
|
fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight)
|
|
bound = 1 / (fan_in**0.5)
|
|
nn.init.uniform_(self.bias, -bound, bound)
|
|
|
|
def forward(self, x: Tensor) -> Tensor:
|
|
return linear(x, self.weight, self.bias)
|