Record Class TensorShardSpec

java.lang.Object
java.lang.Record
smile.llm.parallel.TensorShardSpec
Record Components:
tpSize - tensor-parallel world size.
tpRank - this rank.
numHeads - local query heads (global / tpSize).
numKvHeads - local KV heads.
intermediateSize - local FFN intermediate size.
linearNumKeyHeads - local DeltaNet key heads (may equal global when not sharded).
linearNumValueHeads - local DeltaNet value heads.
public record TensorShardSpec(int tpSize, int tpRank, int numHeads, int numKvHeads, int intermediateSize, int linearNumKeyHeads, int linearNumValueHeads) extends Record
Local shard sizes for a dense decoder under tensor parallelism.

  • Constructor Summary

    Constructors
    Constructor
    Description
    TensorShardSpec(int tpSize, int tpRank, int numHeads, int numKvHeads, int intermediateSize, int linearNumKeyHeads, int linearNumValueHeads)
    Creates an instance of a TensorShardSpec record class.

  • Method Summary

    Modifier and Type
    Method
    Description
    final boolean
    equals(Object o)
    Indicates whether some other object is "equal to" this one.
    static TensorShardSpec
    forRank(int tpSize, int tpRank, int numHeads, int numKvHeads, int intermediateSize, int linearNumKeyHeads, int linearNumValueHeads)
    Builds a per-rank shard spec.
    final int
    hashCode()
    Returns a hash code value for this object.
    int
    intermediateSize()
    Returns the value of the intermediateSize record component.
    int
    linearNumKeyHeads()
    Returns the value of the linearNumKeyHeads record component.
    int
    linearNumValueHeads()
    Returns the value of the linearNumValueHeads record component.
    int
    numHeads()
    Returns the value of the numHeads record component.
    int
    numKvHeads()
    Returns the value of the numKvHeads record component.
    static TensorShardSpec
    single(int numHeads, int numKvHeads, int intermediateSize, int linearNumKeyHeads, int linearNumValueHeads)
    Single-device shard (no tensor parallelism).
    final String
    toString()
    Returns a string representation of this record class.
    int
    tpRank()
    Returns the value of the tpRank record component.
    int
    tpSize()
    Returns the value of the tpSize record component.

    Methods inherited from class Object

    clone, finalize, getClass, notify, notifyAll, wait, wait, wait

  • Constructor Details

    • TensorShardSpec

      public TensorShardSpec(int tpSize, int tpRank, int numHeads, int numKvHeads, int intermediateSize, int linearNumKeyHeads, int linearNumValueHeads)
      Creates an instance of a TensorShardSpec record class.
      Parameters:
      tpSize - the value for the tpSize record component
      tpRank - the value for the tpRank record component
      numHeads - the value for the numHeads record component
      numKvHeads - the value for the numKvHeads record component
      intermediateSize - the value for the intermediateSize record component
      linearNumKeyHeads - the value for the linearNumKeyHeads record component
      linearNumValueHeads - the value for the linearNumValueHeads record component

  • Method Details

    • single

      public static TensorShardSpec single(int numHeads, int numKvHeads, int intermediateSize, int linearNumKeyHeads, int linearNumValueHeads)
      Single-device shard (no tensor parallelism).
      Parameters:
      numHeads - query head count.
      numKvHeads - key/value head count.
      intermediateSize - FFN intermediate size.
      linearNumKeyHeads - DeltaNet key head count.
      linearNumValueHeads - DeltaNet value head count.
      Returns:
      shard with tpSize=1.

    • forRank

      public static TensorShardSpec forRank(int tpSize, int tpRank, int numHeads, int numKvHeads, int intermediateSize, int linearNumKeyHeads, int linearNumValueHeads)
      Builds a per-rank shard spec. Requires numHeads, numKvHeads, and intermediateSize to be divisible by tpSize. Linear-attn head counts must also divide when > 0.
      Parameters:
      tpSize - tensor-parallel size.
      tpRank - this rank.
      numHeads - global query head count.
      numKvHeads - global key/value head count.
      intermediateSize - global FFN intermediate size.
      linearNumKeyHeads - global DeltaNet key heads (0 if unused).
      linearNumValueHeads - global DeltaNet value heads (0 if unused).
      Returns:
      local shard sizes for tpRank.

    • toString

      public final String toString()
      Returns a string representation of this record class. The representation contains the name of the class, followed by the name and value of each of the record components.
      Specified by:
      toString in class Record
      Returns:
      a string representation of this object

    • hashCode

      public final int hashCode()
      Returns a hash code value for this object. The value is derived from the hash code of each of the record components.
      Specified by:
      hashCode in class Record
      Returns:
      a hash code value for this object

    • equals

      public final boolean equals(Object o)
      Indicates whether some other object is "equal to" this one. The objects are equal if the other object is of the same class and if all the record components are equal. All components in this record class are compared with the compare method from their corresponding wrapper classes.
      Specified by:
      equals in class Record
      Parameters:
      o - the object with which to compare
      Returns:
      true if this object is the same as the o argument; false otherwise.

    • tpSize

      public int tpSize()
      Returns the value of the tpSize record component.
      Returns:
      the value of the tpSize record component

    • tpRank

      public int tpRank()
      Returns the value of the tpRank record component.
      Returns:
      the value of the tpRank record component

    • numHeads

      public int numHeads()
      Returns the value of the numHeads record component.
      Returns:
      the value of the numHeads record component

    • numKvHeads

      public int numKvHeads()
      Returns the value of the numKvHeads record component.
      Returns:
      the value of the numKvHeads record component

    • intermediateSize

      public int intermediateSize()
      Returns the value of the intermediateSize record component.
      Returns:
      the value of the intermediateSize record component

    • linearNumKeyHeads

      public int linearNumKeyHeads()
      Returns the value of the linearNumKeyHeads record component.
      Returns:
      the value of the linearNumKeyHeads record component

    • linearNumValueHeads

      public int linearNumValueHeads()
      Returns the value of the linearNumValueHeads record component.
      Returns:
      the value of the linearNumValueHeads record component