Method forward

transformer.py:224–232

x: (B, S, D) kv: (B, S, H, D)

(self,
            x: Tensor,
            kv: Optional[Tensor] = None,)

        return q, k, v

    def forward(self,
            x: Tensor,
            kv: Optional[Tensor] = None,):
        """
        x: (B, S, D)
        kv: (B, S, H, D)
        """
        q, k, v = self._project(x)
        return self._attend(q, k, v, kv_cache=kv)
        

class PreNormAttn(nn.Module):