Update non_xla attention to properly support paged_attention dynamo code path

[wonjoolee95]

• Update non_xla attention to properly support paged_attention dynamo code path
• Fix the original broken dynamo unit tests with paged_attention

Test plan:

root@1fdc3324aeef:/pytorch/xla# python test/test_pallas.py PallasTest.test_paged_attention_wrapper_with_dynamo
.
----------------------------------------------------------------------
Ran 1 test in 1.798s

OK

+ TPU CI

[alanwaketan]

LGTM.

[alanwaketan]

Do you know if this actually initialize anything? I hope not.

[JackCaoG]

yea it is worth checking out, in above we have a warning about the q should be on meta device, I think running ops on meta_tensor will not allocate any device memory.

[wonjoolee95]

According to https://pytorch.org/docs/stable/generated/torch.empty_like.html, it seems like emtpy_like returns an uninitialized tensor. Along with the meta tensor check above, I think we should be good.

[wonjoolee95]

Thanks for the reviews, I'll go ahead and merge this as the CIs (including TPU CI) are all green.