Fix notation and clarification of Puzzle 9

[alexzhang13]

The notation for the softmax in Puzzle 9 is confusing. The current indexing is not representative of the outer product, and also the inclusion of an extra variable B1 is a bit ambiguous. I think the new description (minor change) is more clear.

The new notation also makes it clear the relationship between the k vector and v vector, which is important for understanding how the full flash attention is done.