This document is relevant for: Trn2, Trn3

QKV Batch Shard Kernel API Reference#

QKV batch shard kernel: all_gather on dim 0 + reshape to separate gqa_group_size heads + slice batch.

Implements the Q projection transition from TP64 to TP8DP8 for batch sharding attention. Each rank starts with n Q heads for all B batches, ends with gqa_group_size*n Q heads for B/gqa_group_size batches. The pattern works when either: (1) heads are in dim=0 (NBSd), or (2) n=1 (any layout). It performs all_gather on dim 0, reshapes to (G, dim0, ...), then slices the batch (where G = gqa_group_size):

Input (per rank)         all_gather dim 0       reshape                 slice batch             reshape back
(n, B, S, d)        ->   (G*n, B, S, d)    ->   (G, n, B, S, d)    ->   (G, n, B/G, S, d)  ->   (G*n, B/G, S, d)
(d, B, n=1, S)      ->   (G*d, B, n=1, S)  ->   (G, d, B, n=1, S)  ->   (G, d, B/G, n=1, S)->   (n=G, d, B/G, S)

Example with G=8, n=1, B=32, S=1, d=64:
NBSd: (1,32,1,64) -> (8,32,1,64) -> (8,1,32,1,64) -> (8,1,4,1,64) -> (8,4,1,64)
dBnS: (64,32,1,1) -> (512,32,1,1) -> (8,64,32,1,1) -> (8,64,4,1,1) -> (8,64,4,1)

Background#

The attn_q_batch_shard kernel implements the Q projection transition from a tensor-parallel layout (e.g. TP64) to a tensor-and-data-parallel layout (e.g. TP8DP8) for batch-sharded attention. Each rank starts with n Q heads for all B batches and ends with gqa_group_size * n Q heads for B / gqa_group_size batches.

API Reference#

Source code for this kernel API can be found at: batch_shard.py

attn_q_batch_shard#

nkilib.experimental.collectives.attn_q_batch_shard(input: nl.ndarray, iota_workers: nl.ndarray, gathered_buf: nl.ndarray, gqa_group_size: int, replica_group: ReplicaGroup, layout: AttnQBatchShardLayout = AttnQBatchShardLayout.NBSd, rank_id_in: Optional[nl.ndarray] = None) → nl.ndarray#

QKV batch shard kernel: all_gather on dim 0 + reshape to separate gqa_group_size heads + slice batch.

Parameters:

input (nl.ndarray) – Input Q tensor from TP64 projection. First dim is gathered across gqa_group_size ranks. NBSd: (n_heads, B, S, d) dBnS: (d, B, n_heads, S)
iota_workers (nl.ndarray) – Lookup table mapping rank_id -> batch_offset for scalar_offset DMA. Shape: (1, collective_ranks), values: [(r % gqa_group_size) * B_per_rank for r in range(collective_ranks)] Needed because NKI compiler doesn’t support arithmetic on rank_id.
gathered_buf (nl.ndarray) – Workspace buffer for all_gather result (must be input tensor for scalar_offset)
gqa_group_size (int) – GQA group size (e.g., 8 for TP8DP8, 2 for TP2DP2)
replica_group (ReplicaGroup) – ReplicaGroup defining the collective topology
layout (AttnQBatchShardLayout) – Output layout - NBSd or dBnS
rank_id_in (Optional[nl.ndarray]) – Optional rank_id as input tensor (1,1) int32. If None, uses ncc.rank_id().

Returns:

(gqa_group_size*n_heads, B/gqa_group_size, S, d)

Return type:

nl.ndarray