This document is relevant for: Trn2, Trn3

Attention KV-Parallel Segmented CTE Kernel API Reference#

KV-parallel segmented prefill attention.

Distributes attention computation across ranks, where each rank holds a shard of the KV cache. Uses online softmax to merge partial results.

Background#

The attention_kv_parallel_segmented_cte kernel enables context parallelism for prefill by distributing the KV cache across ranks. Each rank computes attention over its local KV shard in segments, then the partial results are merged across ranks using online softmax. It supports paged (block-based) KV cache, sliding-window masking, and both contiguous and interleaved (round-robin) KV distribution. This kernel replaces the earlier kv_parallel_segmented_prefill kernel.

API Reference#

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

attention_kv_parallel_segmented_cte#

nkilib.core.attention.attention_kv_parallel_segmented_cte(q: nl.ndarray, k_cache: nl.ndarray, v_cache: nl.ndarray, block_tables: nl.ndarray, kvp_q_offset: nl.ndarray, replica_groups: ReplicaGroup, group_size: int, block_size: int, seg_size: int, scale: float = 1.0, global_q_offset: int = 0, tp_out: bool = False, sliding_window: int = 0, kvp_rank_id: Optional[nl.ndarray] = None, kvp_group_size: int = 0) → nl.ndarray#

KV-parallel segmented prefill attention.

Parameters:

  • q (nl.ndarray) – [BS, S, D], This rank’s Q heads (BS = lnc_degree).

  • k_cache (nl.ndarray) – [num_blocks, num_kv_heads, block_size, D], Local KV cache (K).

  • v_cache (nl.ndarray) – [num_blocks, num_kv_heads, block_size, D], Local KV cache (V).

  • block_tables (nl.ndarray) – [1, max_blocks] int32, Block indices for paged KV.

  • kvp_q_offset (nl.ndarray) – [1, 1] int32, Causal mask offset = -rank_id * local_kv_len + global_q_offset. For round-robin KV distribution, set to just global_q_offset since the global-to-local bound conversion handles K-side positioning.

  • replica_groups (ReplicaGroup) – ReplicaGroup for collective operations.

  • group_size (int) – Number of ranks in the replica group.

  • block_size (int) – KV cache block size.

  • seg_size (int) – Segment size for attention iteration.

  • scale (float) – Attention scale factor (default 1.0).

  • global_q_offset (int) – Global token position of Q token 0 (default 0). Used to compute how many prior KV tokens exist within this rank’s shard for each Q chunk.

  • tp_out (bool) – If True, output is transposed to [BS, D, S] (default False).

  • sliding_window (int) – Sliding window size for attention (0 = disabled).

  • kvp_rank_id (Optional[nl.ndarray]) – [1, 1] int32, This rank’s index within the KV-parallel group. Required for interleaved (round-robin) KV distribution to convert global K positions to segment-local positions.

  • kvp_group_size (int) – Number of ranks sharing the KV cache in round-robin fashion. When > 0, enables interleaved KV mode where rank r holds global blocks r, r+R, r+2R, … (R = kvp_group_size).

Returns:

[BS, S, D], Merged attention output for this rank’s Q heads.

Return type:

nl.ndarray

Dimensions:

  • BS: Batch size (lnc_degree = Q heads per physical rank)

  • S: Sequence length

  • D: Head dimension

  • G: Group size (number of ranks per replica group)

This document is relevant for: Trn2, Trn3