Add Tensor Parallel to torch_native_llama

[kwen2501]

Motivation

The torch_native_llama model does not have Tensor Parallel support today. This PR adds it, using torch.distributed APIs.

Modifications

• Added a .tensor_parallel() utility;
• Added ColwiseParallel and RowwiseParallel annotations to related sub-modules;

Tests

pytest test/srt/test_torch_tp.py

Checklist

• Format your code according to the Contributor Guide.
• Add unit tests as outlined in the Contributor Guide.
• Update documentation as needed, including docstrings or example tutorials.

cc: @jerryzh168 @merrymercy @wz337

[jerryzh168]

LGTM, I think this is the best we can do for now until we don't use fused qkv and rely on torch.compile for speedup

[merrymercy]

LGTM. We can merge this soon.

1. Can we add a unit test for this? Example: https://github.com/sgl-project/sglang/blob/main/test/srt/test_torchao.py
2. Can you fix the lint errors: https://sgl-project.github.io/references/contributor_guide.html#format-your-code

[merrymercy]

Add a unit test or at least add some doc strings here on how to run this model?

[kwen2501]

Thanks for the suggestion. I added a section describing how to run the model with TP.

[kwen2501]

@merrymercy Thanks much for your review. Was in roadmapping hence the code change was delayed.
I added a unit test: test/srt/test_torch_tp.py.
It can be triggered by:
pytest test/srt/test_torch_tp.py
and would by default use two GPUs.
It is similar to test/srt/test_bench_latency.py.

Thanks for your rebase and would appreciate your review!