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.NkiTensor, k_cache: nl.NkiTensor, v_cache: nl.NkiTensor, block_tables: nl.NkiTensor, kvp_q_offset: nl.NkiTensor, replica_groups: ncc.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.NkiTensor] = None, kvp_group_size: int = 0, apc_mode: bool = False, valid_num_prior_tokens: nl.NkiTensor = None, fp8_packed: bool = False, k_scale: Optional[nl.NkiTensor] = None, v_scale: Optional[nl.NkiTensor] = None) → nl.NkiTensor#

KV-parallel segmented prefill attention.

Parameters:

  • q (nl.NkiTensor) – [q_heads_per_rank, S, D], This rank’s Q heads.

  • k_cache (nl.NkiTensor) – Local KV cache (K). Shape depends on fp8_packed: - False: [num_blocks, num_kv_heads, block_size, D] - True: [num_blocks, num_kv_heads, block_size // 2, D, 2]

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

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

  • kvp_q_offset (nl.NkiTensor) – [1, 1] int32, Causal mask offset. For contiguous KV sharding: -rank_id * local_kv_len + global_q_offset. For round-robin KV distribution: the actual global Q position (runtime), since the kernel uses kvp_rank_id to handle K-side positioning internally.

  • replica_groups (ncc.ReplicaGroup) – ReplicaGroup for collective operations.

  • group_size (int) – Number of logical ranks in the replica group (= len(replica_groups)).

  • 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) – Compile-time global position of Q token 0 (default 0). In non-APC mode, this is required for correctness — it determines how many prior KV segments are loaded per chunk. In APC mode, set to 0 since the prior token count is derived at runtime from kvp_q_offset instead.

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

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

  • kvp_rank_id (Optional[nl.NkiTensor]) – [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).

  • apc_mode (bool) – Automated prefix caching mode (default False). When True, prior token count is taken from valid_num_prior_tokens (required). Compile-time tile-visibility optimizations are disabled since global_q_offset may not reflect the actual prefix length. Set global_q_offset=0.

  • valid_num_prior_tokens (nl.NkiTensor) – [1, 1] int32. Required when apc_mode is True. The number of fully-visible local prior tokens for Q chunk 0 on this rank. Must be a multiple of block_size. The kernel increments this per chunk by seg_size // kvp_group_size as later chunks see more local KV.

  • fp8_packed (bool) – If True, K uses the packed FP8 layout described above.

  • k_scale (Optional[nl.NkiTensor]) – Optional K-cache dequantization scale, shape [128, 1].

  • v_scale (Optional[nl.NkiTensor]) – Optional V-cache dequantization scale, shape [128, 1]. All entries must repeat one scalar; the kernel applies this tensor along the query partition, not as per-head-dimension values.

Returns:

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

Return type:

nl.NkiTensor

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