support of vectors times matrix in .dot method

[pommedeterresautee]

Currently .dot method do not like when we try a vector times a matrix.
Does MLIR rewrite will help on that?

    assert a.shape[0] >= 16 and a.shape[1] >= 16 and b.shape[1] >= 16, "small blocks not supported!"

[ptillet]

I think it may. We haven't tried yet actually. But note that you can do `sum(A*x[None, :], 1) and this should do a matrix vector product :p

[pommedeterresautee]

tks for the suggestion. I may miss something, but would it be possible to perform a matrix transpose without using .dot args?

[ptillet]

Usually, the way I perform transposition is by constructing ptrs = base + i[:, None] + j[None, :]*stride_j instead of ptrs = base + i[:, None]*stride_i + j[None, :]. You can also just do ptrs = base + i[:, None]*stride_i + j[None, :]*stride_j and the compiler will handle things properly when one of the two strides is equal to 1.

[pommedeterresautee]

tks, sounds obvious after reading it!
Unfortunately still not working in my specific case with chaining mul ops.

Get IndexError: map::at, I hope you will share something obvious I am not yet seeing :-)
Have tried many "variations", even with the TMP workaround hoping something magic would happen... but no.

import torch
import triton
import triton.language as tl


torch.manual_seed(123)


@triton.jit
def kernel(
    Q,
    K,
    V,
    Out,
    q_stride_x,
    k_stride_x,
    k_stride_y,
    v_stride_x,
    v_stride_y,
    out_stride_x,
    SIZE_N: tl.constexpr,
    D_HEAD: tl.constexpr,
):
    size_n_arange = tl.arange(0, SIZE_N)
    d_head_arange = tl.arange(0, D_HEAD)

    # transpose matrix
    k_ptr = K + d_head_arange[:, None] * k_stride_y + size_n_arange[None, :] * k_stride_x
    k = tl.load(k_ptr).to(tl.float32)

    out_ptr = Out + d_head_arange * out_stride_x

    q_ptr = Q + d_head_arange[:, None] * q_stride_x
    q = tl.load(q_ptr).to(tl.float32)

    qk = q * k
    qk = tl.sum(qk, axis=0)[None, :]

    v_ptr = V + size_n_arange[None, :] * v_stride_x + d_head_arange[:, None] * v_stride_y
    v = tl.load(v_ptr).to(tl.float32)

    result = qk * v
    result = tl.sum(result, axis=1)

    tl.store(out_ptr, result)


size_n = 16
d_head = 128

Q = torch.randn((d_head,), dtype=torch.float16, device="cuda")
K = torch.randn((size_n, d_head), dtype=torch.float16, device="cuda")
V = torch.rand_like(K)
out = torch.zeros((d_head,), dtype=torch.float16, device="cuda")

grid = (1,)

kernel[grid](
    Q,
    K,
    V,
    out,
    *Q.stride(),
    *K.stride(),
    *V.stride(),
    *out.stride(),
    size_n,
    d_head,
    num_warps=1,
    num_stages=1,
)


print(f"max error={(out - (Q @ K.t()) @ V).abs().max().item()}")

[pommedeterresautee]

FWIW I made it through 2 calls to the kernel as for loop tends to segfault or refuse to parallelize.
In case reader of this message is interested, code is here: #949

[resorcap]

I encountered the same problem. Promblem solved?

[pommedeterresautee]

The constrain of minimal size still exist but chaining add and mul or using dot with larger tensor than needed and a mask works