Troubles with scaled_dot_product_attention

whatisslove7: <p>Hello, I try to implement my own neural machine translition model with Flash Attention (use scaled_dot_product_attention from torch.nn.functional)<br>
I have two troubles with it:</p>
<ol>
<li>When I wanna use dtype=torch.float16, I have the following error:</li>
</ol>
<pre><code class="lang-auto">RuntimeError: "baddbmm_with_gemm" not implemented for 'Half'
</code></pre>
<ol start="2">
<li>When I try to use device = ‘cuda’ I have this error:</li>
</ol>
<pre><code class="lang-auto">RuntimeError: No available kernel.  Aborting execution.
</code></pre>
<p>For 2 point: I found <a href="https://discuss.pytorch.org/t/using-f-scaled-dot-product-attention-gives-the-error-runtimeerror-no-available-kernel-aborting-execution/180900">this solution</a>, upgrade pytorch for version from post, but it didn’t help me.</p>
<p>Leave here the piece of code with error:</p>
<pre><code class="lang-auto">if self.flash:
        with torch.backends.cuda.sdp_kernel(enable_math=False, enable_flash=True, enable_mem_efficient=False):
            y = F.scaled_dot_product_attention(Q, K, V,
                                              attn_mask=mask,
                                              dropout_p=self.dropout if self.training else 0,
                                              is_causal=True)
        else:
            raise ImportError("PyTorch &gt;= 2.0 must be installed for using Flash Attention")
</code></pre>
<p>I can post the full function code if necessary</p>

ptrblck: <p>Which PyTorch version and which GPU are you using?</p>

whatisslove7: <p>I use GPU T4 (from google colab)<br>
pytorch version: 2.1.0.dev20230618+cu121</p>

ptrblck: <p>Thank you! Could you post the missing code pieces to create a minimal and executable code snippet, please?</p>

whatisslove7: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> hi!</p>
<p>Attention module code:</p>
<pre><code class="lang-auto">class CausalSelfAttention(nn.Module):

    def __init__(self, config):
        super().__init__()

        self.hidden_dim = config.hidden_dim
        self.num_heads = config.num_heads
        self.dropout = config.dropout

        # support is only in PyTorch &gt;= 2.0
        self.flash = hasattr(F, 'scaled_dot_product_attention')

        self.out_linear = torch.nn.Linear(self.hidden_dim, self.hidden_dim)
        self.resid_dropout = nn.Dropout(config.dropout)

    def forward(self, K, V, Q, mask=None):

        batch_size, hidden_dim = Q.size(0), Q.size(2)
        key_len, value_len, query_len = K.size(1), V.size(1), Q.size(1)

        assert hidden_dim % self.num_heads == 0, "Hidden_dim must be equal to num_heads * head_dim"

        K = K.reshape(batch_size, key_len, self.num_heads, -1).transpose(1, 2) # (batch_size, num_heads, seq_len, head_dim)
        V = V.reshape(batch_size, value_len, self.num_heads, -1).transpose(1, 2) # (batch_size, num_heads, seq_len, head_dim)
        Q = Q.reshape(batch_size, query_len, self.num_heads, -1).transpose(1, 2) # (batch_size, num_heads, seq_len, head_dim)

        if self.flash:
            # efficient attention using Flash Attention CUDA kernels
            with torch.backends.cuda.sdp_kernel(enable_math=False, enable_flash=True, enable_mem_efficient=False):
                y = F.scaled_dot_product_attention(Q, K, V,
                                                  attn_mask=mask,
                                                  dropout_p=self.dropout if self.training else 0,
                                                  is_causal=True)
        else:
            raise ImportError("PyTorch &gt;= 2.0 must be installed for using Flash Attention")

        y = y.transpose(1, 2).contiguous().view(batch_size, query_len, hidden_dim)
        # y = self.resid_dropout(self.out_linear(y)) + Q
        return self.resid_dropout(self.out_linear(y))
</code></pre>
<p>Testing code:</p>
<pre><code class="lang-auto">dtype = torch.float16
device = 'cuda'

query = torch.randn(32, 256, 512, device=device)
key = torch.randn(32, 256, 512, device=device)
value = torch.randn(32, 256, 512, device=device)

att = CausalSelfAttention(cfg)
att(key, value, query).shape
</code></pre>
<p>When I set device=device in k, v, q parts, raise RuntimeError error “No available kernel.”<br>
When I set dtype=dtype in k, v, q parts, raise RuntimeError error ““baddbmm_with_gemm” not implemented for ‘Half’”</p>

ptrblck: <p>Your code is not executable as <code>cfg</code> is undefined.</p>

whatisslove7: <p>My mistake<br>
config:</p>
<pre><code class="lang-auto">class config:
    batch_size: int = 16 
    wmt14_size: int = 1E6 
    # model parameters
    hidden_dim: int = 512
    num_heads: int = 8
    dropout: float = 0.1
    # scheduler parameters
    warmup_steps: int = 4000
    #optimizer parameters
    betas: Tuple[float, float] = (0.9, 0.98)
    eps: float = 1.0e-9

cfg = config()
</code></pre>

ptrblck: <p>Thank you!<br>
It seems you are running into known limitation as you are using the <code>SDPA</code> backend manually. E.g. I see:</p>
<pre><code class="lang-python">UserWarning: Expected query, key and value to all be of dtype: {Half, BFloat16}. Got Query dtype: float, Key dtype: float, and Value dtype: float instead. 
</code></pre>
<p>Using an <code>autocast</code> context manager works for me as it’s casting to the expected <code>dtypes</code>:</p>
<pre><code class="lang-python">att = CausalSelfAttention(cfg).to(device)
with torch.cuda.amp.autocast():
    att(key, value, query).shape
</code></pre>
<p>From my understanding the <code>SDPA</code> backends would be selected automatically for you based on some conditions (such as <code>dtype</code>) while the manual implementation could easily fail.</p>

whatisslove7: <p>Thak you so much for your explanation!<br>
But I have one more question<br>
with <code>torch.cuda.amp.autocast()</code>  will there be a selection both the type of <code>SDPA</code> and the type of tensors or only the first one and we need use dtype = bfloat16 in tensors?</p>

ptrblck: <p>By default <code>autocast</code> will use <code>float16</code> and works properly on my device. You can specify the mixed-precision <code>dtype</code> to be <code>bfloat16</code> if you want.</p>

whatisslove7: <p>Thank you so much!<br>
With your help this solution works perfectly</p>