Tensor Parallelism¶

Table of Contents¶

• Concept
• Current Runtime Support
• Weight Sharding Strategies
• Communication Patterns
• Architecture-Aware Shard Plans
• MoE Expert Parallelism
• Embedding and LM Head Sharding
• TP Scheduler and Executor
• Configuration
• Composing with Pipeline Parallelism (2D PP × TP)

Concept¶

Tensor Parallelism (TP) splits individual weight matrices horizontally across GPUs so that each GPU holds a shard of the full weight matrix. All GPUs process the same input tokens, but each computes on a different portion of the weights. Communication collectives (all-reduce, all-gather) are used after each layer to combine the partial results into a complete activation before proceeding to the next layer.

graph LR
    subgraph Input
        X[Input: shape B×S×D]
    end

    subgraph "GPU 0 (Rank 0)"
        W0["W_shard_0<br/>(D × D/2)"]
        P0["Partial output<br/>(B×S×D/2)"]
    end

    subgraph "GPU 1 (Rank 1)"
        W1["W_shard_1<br/>(D × D/2)"]
        P1["Partial output<br/>(B×S×D/2)"]
    end

    X --> W0
    X --> W1
    W0 --> P0
    W1 --> P1
    P0 -->|all-gather| OUT[Full output: B×S×D]
    P1 -->|all-gather| OUT

This is different from pipeline parallelism, where each GPU runs different layers. In TP, every GPU participates in every layer.

Current Runtime Support¶

The currently shipping TP runtime in mlxcel is an in-process multi-rank runtime. A single mlxcel generate or mlxcel serve process loads multiple rank shards and executes them in lockstep inside one process. You do not need --node-role, --peers, or a distributed cluster configuration for this path.

Current tp_size > 1 runtime support:

• Dense Llama-family models
• Dense Qwen2 / Qwen2.5 models
• Dense Qwen3 models
• Dense Qwen3.5 text models
• Gemma 3 text models
• Gemma 4 text models
• ERNIE 4.5 dense models
• Hunyuan v1 Dense models

Current constraints:

• --tp-embedding-mode must stay replicated
• --tp-lm-head-mode must stay replicated
• LoRA adapters are not supported with TP
• TP server batching is enabled for dense Llama, Qwen2/Qwen2.5, Qwen3, Qwen3.5 text, Gemma 3 text, Gemma 4 text, ERNIE 4.5, and Hunyuan v1 Dense
• The TP runtimes decode batched rows in one native batched forward (BatchedDecodeKind::Native), so the server logs decode path: native batched forward for them
• Gemma 4 checkpoints that need the loaded-model fallback (E2B-style, with KV-shared layers) force sequential server execution

Validated local CUDA models:

• models/llama-3.2-1b-4bit
• models/qwen2.5-0.5b-4bit (tp=2, greedy parity through 24 tokens)
• models/qwen2.5-7b-4bit (tp=4)
• models/qwen3-0.6b-4bit
• models/qwen3.5-0.8b-4bit (tp=2)
• models/qwen3.5-4b-4bit (tp=4)
• models/qwen3.5-9b-4bit (tp=2 and tp=4, stepwise CUDA parity)
• models/qwen3.5-27b-4bit (tp=4, stepwise CUDA parity)
• models/gemma3-1b-4bit
• models/ernie-4.5-0.3b-4bit
• models/hunyuan-1.8b-4bit

Not currently enabled:

• TP + LoRA

Weight Sharding Strategies¶

Defined in src/distributed/tensor_parallel/shard_strategy.rs:

Column Parallel¶

Splits the output dimension (columns) of a weight matrix across ranks. Each rank computes a subset of the output features.

block-beta
    columns 3
    W["W: [D_in × D_out]"]:3
    R0["Rank 0\nW[:, 0 : D_out/N]"] R1["Rank 1\nW[:, D_out/N : 2·D_out/N]"] Rn["...\nRank N-1"]

Used for: attention Q/K/V projections, FFN gate and up projections.

After the column-parallel operation, an all-gather is required if the full output is needed immediately, or a reduce-scatter can be deferred if the next operation is row-parallel.

Row Parallel¶

Splits the input dimension (rows) of a weight matrix across ranks. Each rank holds a partial activation that contributes to a partial sum, which is combined via all-reduce.

block-beta
    columns 3
    W["W: [D_in × D_out]"]:3
    R0["Rank 0\nW[0 : D_in/N, :]"] R1["Rank 1\nW[D_in/N : 2·D_in/N, :]"] Rn["...\nRank N-1"]

Used for: attention output projection, FFN down projection.

After the row-parallel matmul, an all-reduce (sum) is applied to produce the correct combined output.

Vocab Parallel¶

Splits the vocabulary dimension of the embedding table and LM head across ranks. Each rank is responsible for a contiguous slice of the token vocabulary.

Used for: embedding lookup, LM head output logits.

Replicated¶

The weight is fully replicated on every rank. No sharding or communication required. Used for small tensors such as layer normalization parameters and bias vectors.

Communication Patterns¶

Defined in src/distributed/tensor_parallel/collective.rs:

graph TD
    subgraph AllReduce["All-Reduce (Ring)"]
        AR0[Rank 0<br/>partial sum] -->|Reduce-scatter| RS[Each rank<br/>holds one reduced chunk]
        AR1[Rank 1<br/>partial sum] --> RS
        AR2[Rank 2<br/>partial sum] --> RS
        RS -->|All-gather| AG[All ranks<br/>hold full reduced tensor]
    end

    subgraph AllGather["All-Gather"]
        AG0[Rank 0 shard] -->|ring forwarding| FG[All ranks<br/>hold full tensor]
        AG1[Rank 1 shard] --> FG
    end

    subgraph ReduceScatter["Reduce-Scatter"]
        RS0[Rank 0 full tensor] -->|accumulate + scatter| RSO[Each rank<br/>holds 1/N reduced shard]
        RS1[Rank 1 full tensor] --> RSO
    end

All-reduce (ring):

The most common collective in TP. Used after row-parallel matrix multiplications to sum partial results across all ranks. Uses a two-step ring algorithm: reduce-scatter followed by all-gather. Total data transfer per rank: 2 × (N-1)/N × tensor_size.

All-gather:

Used to assemble sharded outputs into a complete tensor. Each rank contributes its shard; all ranks end up with the full tensor.

Reduce-scatter:

Used in optimized pipelines where the result of a row-parallel operation is immediately consumed by a column-parallel operation. Avoids the all-gather step when only a shard of the output is needed.

All collectives use a ring topology. Rank r sends to (r+1) % N and receives from (r+N-1) % N in each step. Over N-1 steps, every rank has communicated with every other rank exactly once.

Architecture-Aware Shard Plans¶

mlxcel generates architecture-specific shard plans via src/distributed/tensor_parallel/plan_generator.rs. The plan encodes which weight tensor gets which ShardStrategy and what CommPattern follows.

Standard Transformer (Llama, Qwen, Gemma, etc.)¶

GQA-Aware Attention Sharding¶

For models with Grouped Query Attention (GQA), the number of KV heads is often much smaller than the number of Q heads. The shard plan accounts for this:

• If num_kv_heads >= tp_size: K/V projections are sharded normally.
• If num_kv_heads < tp_size: K/V projections are replicated on all ranks to avoid fragmenting KV heads smaller than one per rank.

This is computed in src/distributed/tensor_parallel/parallel_attention.rs.

graph LR
    subgraph "GQA Sharding (tp_size=4, 32 Q heads, 8 KV heads)"
        Q["Q heads: 32<br/>2 per rank"]
        K["K heads: 8<br/>2 per rank"]
        V["V heads: 8<br/>2 per rank"]
    end

    subgraph "GQA Replication (tp_size=8, 4 KV heads)"
        Q2["Q heads: 32<br/>4 per rank"]
        K2["K heads: 4<br/>replicated (< 1 per rank)"]
        V2["V heads: 4<br/>replicated (< 1 per rank)"]
    end

SSM/Mamba¶

For Mamba and Mamba2 models, the convolution and SSM state matrices are sharded along their channel dimension, following the same column/row parallel pattern.

MoE Expert Parallelism¶

The following modes describe the shard-plan design. They are not enabled by the current production tp_size > 1 runtime, which is restricted to the dense families listed above.

For Mixture-of-Experts models (Mixtral, DeepSeek V2/V3, Qwen MoE, etc.), mlxcel supports two MoE sharding modes, configurable via --tp-moe-mode:

Expert Parallel (default)¶

Whole experts are assigned to ranks round-robin: expert i lives entirely on rank i % tp_size. All expert weights for a given rank are local: no intra-expert communication is needed.

flowchart LR
    subgraph "Rank 0"
        E0["Expert 0"]
        E2["Expert 2"]
    end
    subgraph "Rank 1"
        E1["Expert 1"]
        E3["Expert 3"]
    end
    Router --> E0 & E1 & E2 & E3

Best when num_experts >= tp_size. The router still runs on all ranks, but each rank only executes its locally-assigned experts.

Within-Expert Parallel¶

Each expert's internal FFN weights are sharded across all ranks using the same column-parallel gate/up and row-parallel down projection pattern as a regular dense FFN.

Best when num_experts < tp_size (unusual but possible for small expert counts with many GPUs).

Defined in src/distributed/tensor_parallel/parallel_moe.rs.

Embedding and LM Head Sharding¶

The configuration types model these strategies, but the runtime requires both embedding and LM-head mode to remain replicated; the sharded alternatives below are not production-enabled.

Two modes, configurable via --tp-embedding-mode and --tp-lm-head-mode:

Vocab-parallel mode is recommended for models with very large vocabularies (> 100K tokens) when memory is the primary constraint.

TP Scheduler and Executor¶

TP execution is orchestrated by two components:

TP Scheduler (rank 0 only)¶

src/distributed/tensor_parallel/tp_scheduler.rs

The scheduler runs exclusively on rank 0. It is responsible for:

• Admitting new sequences and managing the prefill queue.
• Composing batches (selecting which sequences to prefill or decode per step).
• Broadcasting StepDecision messages to all other ranks before each inference step.
• Coordinating token sampling: rank 0 samples the next token and broadcasts the result.
• Continuous batching: interleaving prefill and decode steps for maximum GPU utilization.

No other rank makes independent scheduling decisions. All ranks wait for the scheduler's broadcast before executing each step.

Sequence lifecycle:

stateDiagram-v2
    [*] --> Waiting: request submitted
    Waiting --> PendingPrefill: admitted to batch
    PendingPrefill --> Decoding: prefill complete
    Decoding --> Completed: EOS or max_tokens
    Decoding --> Evicted: preempted
    Evicted --> Waiting: re-admitted
    Completed --> [*]

TP Executor (all ranks)¶

src/distributed/tensor_parallel/tp_executor.rs

Every rank (including rank 0) runs a TPExecutor. The executor:

• Receives StepDecision values from the rank 0 broadcast.
• Maintains identical sequence state on every rank (for lockstep execution).
• Participates in barrier synchronization before each step.
• On decode steps, rank 0 samples the next token; other ranks receive the sampled token via broadcast.
• Reports rank health status; errors are propagated to stop the entire TP group.

sequenceDiagram
    participant S as Scheduler (rank 0)
    participant E0 as Executor rank 0
    participant E1 as Executor rank 1
    participant E2 as Executor rank 2

    S->>E0: broadcast StepDecision
    S->>E1: broadcast StepDecision
    S->>E2: broadcast StepDecision
    Note over E0,E2: All ranks execute forward pass in lockstep
    E0->>E1: all-reduce (per layer)
    E1->>E2: all-reduce (per layer)
    E2->>E0: all-reduce (per layer)
    Note over E0: Rank 0 samples next token
    E0->>E1: broadcast sampled token
    E0->>E2: broadcast sampled token
    Note over S: Update sequence state

Configuration¶

For the current runtime, leave both embedding-related flags at replicated.

Example: local 2-rank TP text generation on one machine:

mlxcel generate \
  --model ./models/llama-3.2-1b-4bit \
  --prompt "Explain tensor parallelism in one sentence." \
  --tp-size 2

Example: local 2-rank TP server:

mlxcel serve \
  --model ./models/qwen3-0.6b-4bit \
  --tp-size 2

Composing with Pipeline Parallelism (2D PP × TP)¶

Tensor Parallelism can be composed with Pipeline Parallelism to form a 2D (pp_stage, tp_rank) mesh. Use this when a single stage of a PP-only topology still exceeds one node's memory. The typical example is a Llama-70B-class model served on four medium-sized GPUs as a 2 PP stages × 2 TP ranks mesh.

See also: Pipeline Parallelism → Composing with Tensor Parallelism (2D PP × TP) for the PP-side perspective on the same mesh.

How TP behaves inside a stage¶

Within a single pipeline stage, the TP collective layer is unchanged from the 1D description above: rank-0 drives admission and broadcasts step decisions, and every rank on the stage participates in the ring-topology all-reduce / reduce-scatter / all-gather for that stage's layers. TP peer discovery is stage-local: all_reduce_sum, all_gather, and reduce_scatter only talk to the other TP ranks that share the same pp_stage.

How PP boundaries work across TP ranks¶

At a PP stage boundary, each tp_rank hands its activation shard to the node at the next stage that holds the same tp_rank. That is, stage 0 rank 0 talks to stage 1 rank 0; stage 0 rank 1 talks to stage 1 rank 1. There is no cross-rank exchange at the boundary because the TP shard layout is uniform across stages.

Transport traffic classes¶

The 2D routing layer shares a single connection pool across every node in the mesh, but it distinguishes two traffic classes so that the right handler processes each message:

Both classes are selected via PipelineTensorParallelRouter; the registry lookups used by the collective layer (NodeRegistry::nodes_at_stage) and the activation-transfer layer (NodeRegistry::nodes_at_rank) provide the full peer set for each class.

CLI entry point¶

mlxcel generate \
  --model ./models/llama-70b-4bit \
  --pp-size 2 \
  --tp-size 2 \
  -p "Hello" -n 50

The validator in src/commands/generate.rs accepts the combination and emits a readable error on inconsistent topologies.

Cluster TOML and cache coordination¶

The cluster TOML uses the pipeline_tensor_parallel role (aliases: pp_tp, 2d) and requires exactly pp_size × tp_size such nodes, each with both stage and rank. See pipeline-parallelism.md → Composing with Tensor Parallelism (2D PP × TP) for a full 2×2 TOML example and the cache coordination rules (coordinated_2d_admission, broadcast_2d_eviction).

Limitations¶

• Uniform TP width per stage: heterogeneous widths (e.g., stage 0 at tp=4, stage 1 at tp=2) are a future extension.
• No 3D parallelism: expert-parallelism (PP × TP × EP) is tracked separately.