Class QuantTpSharding

java.lang.Object
smile.llm.quant.QuantTpSharding
public final class QuantTpSharding extends Object
Shard-then-pack helpers for tensor-parallel quantized linears.

Column-parallel shards out_features; row-parallel shards in_features. Callers pack the local shard into Marlin / FP8 after sharding (never pack globally then slice).

  • Method Details

    • shardColumn

      public static Tensor shardColumn(Tensor weight, int tpSize, int tpRank)
      Column-parallel shard of a dense weight [out, in] (or FP8).
      Parameters:
      weight - global weight.
      tpSize - TP world size.
      tpRank - this rank.
      Returns:
      local [out/tp, in] (caller owns).

    • shardRow

      public static Tensor shardRow(Tensor weight, int tpSize, int tpRank)
      Row-parallel shard of a dense weight [out, in].
      Parameters:
      weight - global weight.
      tpSize - TP world size.
      tpRank - this rank.
      Returns:
      local [out, in/tp] (caller owns).

    • shardGptqQweightColumn

      public static Tensor shardGptqQweightColumn(Tensor qweight, int tpSize, int tpRank)
      Column-parallel shard of GPTQ qweight packed along input dim0 ([in/8, out]). Shards along out.
      Parameters:
      qweight - GPTQ packed weights.
      tpSize - tensor-parallel world size.
      tpRank - this rank.
      Returns:
      sharded qweight (caller owns).

    • shardGptqQweightRow

      public static Tensor shardGptqQweightRow(Tensor qweight, int tpSize, int tpRank)
      Row-parallel shard of GPTQ qweight [in/8, out] along packed K.
      Parameters:
      qweight - GPTQ packed weights.
      tpSize - tensor-parallel world size.
      tpRank - this rank.
      Returns:
      sharded qweight (caller owns).

    • shardAwqQweightColumn

      public static Tensor shardAwqQweightColumn(Tensor qweight, int tpSize, int tpRank)
      Column-parallel shard of AWQ qweight [in, out/8] along packed out.
      Parameters:
      qweight - AWQ packed weights.
      tpSize - tensor-parallel world size.
      tpRank - this rank.
      Returns:
      sharded qweight (caller owns).

    • shardAwqQweightRow

      public static Tensor shardAwqQweightRow(Tensor qweight, int tpSize, int tpRank)
      Row-parallel shard of AWQ qweight [in, out/8] along K.
      Parameters:
      qweight - AWQ packed weights.
      tpSize - tensor-parallel world size.
      tpRank - this rank.
      Returns:
      sharded qweight (caller owns).

    • shardScalesColumn

      public static Tensor shardScalesColumn(Tensor scales, int tpSize, int tpRank)
      Column-parallel shard of scales [groups, out] along out.
      Parameters:
      scales - group scales.
      tpSize - tensor-parallel world size.
      tpRank - this rank.
      Returns:
      sharded scales (caller owns).

    • shardScalesRow

      public static Tensor shardScalesRow(Tensor scales, int tpSize, int tpRank)
      Row-parallel shard of scales [groups, out] along groups (K groups).
      Parameters:
      scales - group scales.
      tpSize - tensor-parallel world size.
      tpRank - this rank.
      Returns:
      sharded scales (caller owns).