Add sdpa arg comments (#5323)

Summary: Pull Request resolved: #5323

Reviewed By: JacobSzwejbka

Differential Revision: D62623249

Pulled By: dvorjackz

fbshipit-source-id: 468abd913a4dcb9b2474ec34881cfcec2654a024

Author: jackzhxng

examples/models/llama2/llama_transformer.py

@@ -218,9 +218,9 @@ def __init__(
     def forward(
         self,
         input_pos: torch.Tensor,
-        q: torch.Tensor,
-        k: torch.Tensor,
-        v: torch.Tensor,
+        q: torch.Tensor,  # Already have rotary embeddings. (bs, seqlen, n_local_heads, head_dim)
+        k: torch.Tensor,  # Already have rotary embeddings. (bs, seqlen, n_local_kv_heads, head_dim)
+        v: torch.Tensor,  # (bs, seqlen, n_local_kv_heads, head_dim)
         bsz,
         seqlen,
         mask: torch.Tensor,