ggml_tp_dp_forward: TPxDP hybrid matrix multiply across replicas of Vulkan device groups

Description

Computes Y = X %*% t(W) as a hybrid of tensor parallelism and data parallelism: the weight matrix W is replicated across replicas device groups (data parallelism over the batch X), and within each group the weight rows are split across that group's GPUs (tensor parallelism, via ggml_vulkan_split_mul_mat). Each replica computes a contiguous shard of the batch rows; the shards are concatenated back into the full result.

Usage

ggml_tp_dp_forward(
  W,
  X,
  replicas,
  weights = NULL,
  transport = c("host-staging", "opaque-fd", "device-group")
)

Value

The M x N result matrix, equal to X %*% t(W) up to the GPU's floating-point accumulation.

Arguments

W: Weight matrix, N x K (replicated to every group).
X: Activation matrix, M x K; its rows (the batch) are split across the replicas.
replicas: A list of integer vectors, each a group of physical GPU indices (0-based) forming one replica, e.g. list(c(0, 1), c(2, 3)). Within a group the weights are tensor-split; across groups the batch is data-split.
weights: Optional numeric vector (length = group size) giving the per-device TP row weighting inside each group (applied to every group). NULL (default) splits evenly.
transport: Cross-device TP gather transport (see ggml_vulkan_split_mul_mat).

Details

This mirrors the inference layout for the 4x P100 target (2 replicas x TP=2): there is no cross-replica communication, so throughput scales with the replica count while each replica's cross-device TP gather stays local to its group.

Examples

Run this code

# \donttest{
if (ggml_vulkan_available() && ggml_vulkan_device_count() >= 4) {
  W <- matrix(rnorm(2048 * 64), nrow = 2048)
  X <- matrix(rnorm(8 * 64), nrow = 8)      # batch of 8
  # 2 replicas x TP=2: batch split 4/4, weights split within {0,1} and {2,3}
  Y <- ggml_tp_dp_forward(W, X, replicas = list(c(0, 1), c(2, 3)))
  max(abs(Y - X %*% t(W)))
}
# }