Scaled_dot_product_attention not useful for inference

mseeger: <p>I am implementing long context inference and noticed an odd deficiency of <code>scaled_dot_product_attention</code>.</p>
<p>In long context inference, a prompt is processed in chunks, so that <code>query</code> is shorter than <code>key</code>, <code>value</code> (in the doc: <code>L &lt; S</code>). I can pass <code>attn_mask</code> and set <code>is_causal=False</code>, but noted that in this case, often only the naive C++ implementation is supported. I’d therefore like to use the <code>is_causal=True</code> case with <code>L &lt; S</code>.</p>
<p>But the way this is implemented, it is not really useful. If <code>i</code> indexes <code>query</code>, <code>j</code> indexes <code>key</code>, the masking is such that what is computed is <code>inner(q_i, k_j) + (-infty) * ind(i &lt; j)</code>. This means that most entries are <code>-infty</code>. Say, <code>L = 3, S = 5</code>. The mask matrix is:</p>
<pre><code class="lang-auto">tensor([[0., -inf, -inf, -inf, -inf],
        [0., 0., -inf, -inf, -inf],
        [0., 0., 0., -inf, -inf]])
</code></pre>
<p>That makes no sense, why would I ever want to mask out most score entries? What is needed for inference, is that the mask is</p>
<pre><code class="lang-auto">tensor([[0., 0., 0., -inf, -inf],
        [0., 0., 0., 0., -inf],
        [0., 0., 0., 0., 0.]])
</code></pre>
<p>This would correspond to a new chunk of size 3, whereas the KV cache has length 5 (the final 3 slots are for the new tokens).</p>
<p>I do wonder what is the point of this definition of the mask if <code>is_causal=True</code>, given that only in this special case (so <code>attn_mask=None</code>), the most efficient SDPA implementations are given.</p>
<p>Sure, I can just use my own implementation in this case, or send in <code>attn_mask</code> (in which case the efficient SDPA implementations cannot be used), but why?</p>
<p>Note that for large context inference, this IS important. I might have a KV cache of size (say) <code>S = 2 ** 16</code> and a blocksize for processing the prompt of size <code>L = 2 ** 10</code>, in which case I’d appreciate a fast SDPA implementation.</p>
<p>As it stands, the case <code>L &lt; S</code> and <code>is_causal=True</code> seems not useful to me for anything I’d want to do with SDPA. But if the default mask in this case was defined just a bit differently (see above, more 0s, less -infty), then that would be really useful for long context inference.</p>

mseeger: <p>Having worked quite a bit more on this, I found it quite striking while the “training mode” (or <code>is_causal=True</code>) is really well covered by a number of different efficient SDPA kernels, the situation for <code>is_causal=False</code> is pretty uncovered. Or at least the offering does not make so much sense.</p>
<p>Do you guys notice that by far the most time in long context inference is spent on the chunks after the first prefill one? Where the rest of the prompt is processed in chunks and a KV cache is served?</p>
<p>Say I have a prompt of 150000 tokens. On GPUs I can afford, no more than 32k of that can be done by the first “prefill” step. If the remaining sequence is split into chunks of size 1k, I have about 120 of them to process sequentially. And each step needs SDPA.</p>
<p>And when I call your SDPA kernels for any of these, I just see OOM immediately. Mainly because of the very simple reason that the <code>mask</code> argument you require is already too large.</p>
<p>I had to implement all of this from scratch, including splitting the tensors into shards, and generating the masks on the fly. I have no idea about Cuda, so this could likely be faster. But at least, it works. Fair enough, but guess how many folks need to run efficient inference with long contexts, vs how many need to do training …</p>