Incomplete implementation of SparseGEMV

[kyang-06]

TEAL/kernels/sparse_gemv.py

Line 271 in fb7373c

return splitk_sparse_gemv(hidden_states, weights, threshold, sparsity_bin) if hidden_states.shape[1] == 1 else torch.matmul(hidden_states, weights.T)

Hi, I notice that the SparseGEMV kernel only manage the case when batch_size=1 & seqlen=1. Beyond that case, the kernel outputs wrong answer.

Is it expected that this kernel only work for decoding stage? Then where is the implementation about Appendix A4?

[nil0x9]

They mentioned in A4 that single-batch setting is used. That said, I don't think it's appropriate to compare 2:4 sparsity here as 2:4 sparsity is not fit for small-batch matmul's.

[quaternior]

I think this method(therefore the kernel) only focus on single-batch setting in decoding stage, as other activation sparsity also only focus on them. The advantage of activation sparsity is latency, not a throughput I think, and Appendix A4 might be focused on latency(as speedup)

[chromecast56]

yes its primarily a decoding kernel. A.4 is based on loss evals, where the sparsity (both weight and activation) is simulated