Enable attn_weight scalar

Author: alanwaketan

src/transformers/models/llama/modeling_llama.py

@@ -381,8 +381,8 @@ def forward(
  attn_output = torch.matmul(attn_weights, value_states)
  else:
  # Integrated with PyTorch/XLA Pallas Flash Attention:
- # TODO: enable 1 / math.sqrt(self.head_dim).
  from torch_xla.experimental.custom_kernel import flash_attention
+ query_states = query_states / math.sqrt(self.head_dim)
  attn_output = flash_attention(query_states, key_states, value_states, causal=True, partition_spec=('fsdp', None, None, None))
 
  if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):