Make torch TP composable with torchao

[kwen2501]

Motivation

The TP styles in PyTorch trunk performs a scatter from rank 0 to distribute a full tensor into DTensors.
While this is safer for training, it is unnecessary for inference where all ranks' weight already come from the same checkpoint. Therefore, we choose to do a local sharding instead of scatter.

Modifications

This PR customizes the TP styles impls to use _shard_tensor instead of distribute_tensor.

Worked configs:

• TP + int8wo
• TP + fp8wo

$ export ENABLE_INTRA_NODE_COMM=1
$ python3 -m sglang.bench_one_batch --model meta-llama/Meta-Llama-3-8B --batch-size 1 --input 128 --output 8 --json-model-override-args '{"architectures": ["TorchNativeLlamaForCausalLM"]}' --enable-torch-compile --torchao-config int8wo --tp 4

Output:

Benchmark ...
Prefill. latency: 0.05945 s, throughput:   2153.25 token/s
Decode.  latency: 0.00321 s, throughput:    311.98 token/s
Decode.  latency: 0.00299 s, throughput:    334.26 token/s
Decode.  latency: 0.00298 s, throughput:    335.68 token/s
Decode.  latency: 0.00297 s, throughput:    337.05 token/s
Decode.  latency: 0.00296 s, throughput:    337.33 token/s
Decode.  median latency: 0.00297 s, median throughput:    337.05 token/s
Total. latency:  0.080 s, throughput:   1690.04 token/s

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.