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:
0z5a
2026-09-02 13:11:25 +08:00
parent 9c3ef0c2a1
commit a144d7f306
14 changed files with 1242 additions and 3 deletions
+3 -2
View File
@@ -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)