x: tensor of shape (M) Y: matrix of shape (M, B) return: matrix-vector product of Y and x, shaped (B)
(
x: torch.Tensor,
Y: torch.Tensor, # noqa: N803
)
| 4 | |
| 5 | |
| 6 | def tl_gemv( |
| 7 | x: torch.Tensor, |
| 8 | Y: torch.Tensor, # noqa: N803 |
| 9 | ) -> torch.Tensor: |
| 10 | """ |
| 11 | x: tensor of shape (M) |
| 12 | Y: matrix of shape (M, B) |
| 13 | return: matrix-vector product of Y and x, shaped (B) |
| 14 | """ |
| 15 | assert x.shape[0] == Y.shape[0], (x.shape, Y.shape) |
| 16 | M, B = Y.shape # noqa: N806 |
| 17 | out = torch.empty(B, device=x.device, dtype=x.dtype) |
| 18 | |
| 19 | def grid(meta): |
| 20 | return (triton.cdiv(B, meta["BLOCK_SIZE_B"]),) |
| 21 | |
| 22 | tl_gemv_kernel[grid](x, Y, out, M, B) |
| 23 | return out |
| 24 | |
| 25 | |
| 26 | @triton.autotune( |