Implementation:Mlc ai Mlc llm Tensor Parallel

Overview

Tensor Parallel is a support module in MLC LLM that defines sharding operators used for tensor parallelism in large language model inference. It is located at python/mlc_llm/support/tensor_parallel.py (116 lines).

The module provides the ShardSingleDim dataclass, which describes how to shard a tensor along a single dimension, and a shard_bias context manager for handling bias terms during tensor-parallel execution. These components work with TVM's Tensor Expression (TE) and Tensor IR (TIR) infrastructure to generate efficient sharding kernels.

Purpose

When distributing a model across multiple devices for tensor parallelism, weight matrices must be split (sharded) along specific dimensions. For example, an attention layer's query projection weight may be split along the output dimension so each device handles a subset of attention heads. This module defines the sharding strategy and generates the TIR compute functions that perform the actual tensor splitting.

Classes

ShardSingleDim

@dataclasses.dataclass
class ShardSingleDim:
    name: str
    dim: int
    segs: Optional[List[int]] = None

A dataclass describing a single-dimension sharding strategy.

Parameters:

• name -- The name of the shard function, used for identification in the IR module.
• dim -- The dimension along which to shard the tensor.
• segs -- Optional list of segment lengths along the shard dimension. When None, the entire dimension is treated as a single segment. When specified, the dimension is divided into segments of potentially different lengths, each sharded independently and evenly across workers.

The segs parameter is useful for fused weight matrices where, for example, query, key, and value projections are concatenated along a single dimension but have different sizes.

gen_tir

def gen_tir(self, shards: int, weight: nn.Tensor) -> tir.PrimFunc:

Generates a TIR PrimFunc that shards the weight tensor. The algorithm:

1. Computes the input shape by scaling the shard dimension: shape[dim] * shards.
2. Creates a TE placeholder for the full (unsharded) weight.
3. For each segment along the shard dimension:
   • Slices the relevant portion of the input tensor.
   • Reshapes to insert a new shards dimension: [..., shards, seg_size, ...].
   • Transposes to move the shard dimension to the front: [shards, ..., seg_size, ...].
4. Concatenates all segments along the shard dimension (axis 1 + dim).
5. Returns a tir.PrimFunc that maps the input tensor to the sharded output.

The output has shape [num_shards, *original_weight_shape].

All shape values are cast to int64 to prevent integer overflow:

shape = [tir.IntImm("int64", v) for v in shape]
segs = [tir.IntImm("int64", v) for v in segs]

gen_shard_info

def gen_shard_info(self, shards: int, weight: nn.Tensor) -> Dict[str, Any]:

Generates metadata about the sharding operation, including:

• func_name -- The shard function name
• in_shape -- The full (unsharded) input shape
• out_shape -- The sharded output shape (shards, *weight.shape)
• out_dtype -- The weight data type

_compute_in_shape

def _compute_in_shape(self, shards: int, weight: nn.Tensor) -> List[int]:

Computes the full weight shape before sharding by scaling the shard dimension:

return [*shape[:self.dim], shape[self.dim] * shards, *shape[self.dim + 1:]]

Context Manager

shard_bias

@contextmanager
def shard_bias(linear: nn.Linear, tensor_parallel_shards: int):

A context manager that temporarily divides a linear layer's bias by the number of tensor parallel shards. This is necessary because when the outputs of tensor-parallel linear layers are summed (via all-reduce), the bias would be added multiple times. Dividing by the shard count ensures the correct bias value after aggregation.

original_bias = linear.bias
if tensor_parallel_shards > 1:
    linear.bias = linear.bias / tensor_parallel_shards
yield
linear.bias = original_bias

The context manager restores the original bias value upon exit, making it safe for repeated use.

Relationship to Preshard

The ShardSingleDim class is used as the shard_strategy attribute on model parameters. The Preshard module reads this attribute and calls gen_tir to generate the actual sharding IR functions during weight conversion.

Dependencies

• dataclasses -- Python standard library for the dataclass decorator
• contextlib.contextmanager -- For the shard_bias context manager
• tvm.te -- TVM Tensor Expression for compute definitions
• tvm.tir -- TVM Tensor IR for low-level function representation
• tvm.topi -- TVM operator inventory (transpose, reshape, concatenate)
• tvm.relax.frontend.nn -- Neural network abstractions (nn.Parameter, nn.Linear, nn.Tensor)

File Location

python/mlc_llm/support/tensor_parallel.py