Class ColumnParallelLinear

java.lang.Object

smile.llm.parallel.ColumnParallelLinear

public final class ColumnParallelLinear
extends Object

Megatron-style column-parallel linear: shards out_features across TP ranks. Input is replicated; output is the local shard (no gather).

Supports dense LinearLayer or quantized LinearOp (FP8 / NVFP4 / Marlin). Quantized weights must be sharded then packed before constructing this wrapper.

Constructor Summary

Constructors

Constructor	Description
ColumnParallelLinear(int inFeatures, int globalOutFeatures, boolean bias, int tpSize, int tpRank)	Creates a dense column-parallel linear layer for the given TP rank.
ColumnParallelLinear(LinearOp local, int globalOutFeatures, int tpSize, int tpRank)	Wraps an already-sharded (and packed, if quantized) local linear op.

Method Summary

Modifier and Type	Method	Description
Tensor	forward(Tensor input)	Applies the local linear transform to input.
LinearLayer	linear()	Returns the underlying dense linear layer.
LinearOp	linearOp()	Returns the local LinearOp (dense or quantized).
int	localOutFeatures()	Returns the local output feature count for this TP rank.
MemorySegment	module()	Returns the native module handle for weight registration (dense only).
int	tpRank()	Returns this tensor-parallel rank.
int	tpSize()	Returns the tensor-parallel world size.

Methods inherited from class Object

clone, equals, finalize, getClass, hashCode, notify, notifyAll, toString, wait, wait, wait

Constructor Details

ColumnParallelLinear

public ColumnParallelLinear(int inFeatures,
 int globalOutFeatures,
 boolean bias,
 int tpSize,
 int tpRank)

Creates a dense column-parallel linear layer for the given TP rank.

Parameters:

inFeatures - input feature count.

globalOutFeatures - global output feature count (must be divisible by tpSize).

bias - whether to include a bias term.

tpSize - tensor-parallel world size.

tpRank - this rank in [0, tpSize).

ColumnParallelLinear

public ColumnParallelLinear(LinearOp local,
 int globalOutFeatures,
 int tpSize,
 int tpRank)

Wraps an already-sharded (and packed, if quantized) local linear op.

Parameters:

local - already-sharded local linear op.

globalOutFeatures - global output feature count (must be divisible by tpSize).

tpSize - tensor-parallel world size.

tpRank - this rank in [0, tpSize).

Method Details

linearOp

public LinearOp linearOp()

Returns the local LinearOp (dense or quantized).

Returns:

local LinearOp (dense or quantized).

linear

public LinearLayer linear()

Returns the underlying dense linear layer.

Returns:

dense LinearLayer.

Throws:

IllegalStateException - if this wrapper holds a quantized op.

module

public MemorySegment module()

Returns the native module handle for weight registration (dense only).

Returns:

native module handle.

localOutFeatures

public int localOutFeatures()

Returns the local output feature count for this TP rank.

Returns:

globalOutFeatures / tpSize.

tpRank

public int tpRank()

Returns this tensor-parallel rank.

Returns:

rank in [0, tpSize).

tpSize

public int tpSize()

Returns the tensor-parallel world size.

Returns:

TP world size.

forward

public Tensor forward(Tensor input)

Applies the local linear transform to input.

Parameters:

input - input tensor.

Returns:

local output shard.