This document is relevant for: Trn2, Trn3

MS Deformable Attention Kernel API Reference#

Multi-scale deformable attention kernel that uses indirect DMA transpose.

Background#

The ms_deformable_attention kernel computes multi-scale deformable attention, sampling values from multiple feature-pyramid levels at learned sampling locations using bilinear interpolation. It uses indirect DMA transpose to gather the sampled values efficiently and supports both BLNC/BNLC value layouts and BQHLP2/B2QHLP sampling-location layouts.

API Reference#

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

ms_deformable_attention#

nkilib.experimental.deformable_attention.ms_deformable_attention(value: nl.ndarray, spatial_shapes: tuple, level_start_index: tuple, sampling_locations: nl.ndarray, attention_weights: nl.ndarray, value_layout: str = 'BLNC', sampling_locations_layout: str = 'BQHLP2', align_corners: bool = False, padding_mode: str = 'zeros') → nl.ndarray#

Multi-scale deformable attention kernel that uses indirect DMA transpose.

Parameters:

  • value (nl.ndarray) – Value tensor in HBM. Shape depends on value_layout: - If value_layout=”BLNC”: (B, L, N_h, C_h) - If value_layout=”BNLC”: (B, N_h, L, C_h)

  • spatial_shapes (tuple) – Tuple of (H_i, W_i) tuples specifying spatial dimensions for each level

  • level_start_index (tuple) – Tuple of start indices for each level in the flattened L dimension

  • sampling_locations (nl.ndarray) – Normalized sampling coordinates in HBM. Shape depends on layout: - If sampling_locations_layout=”BQHLP2”: (B, N_q, N_h, N_l, N_p, 2) - If sampling_locations_layout=”B2QHLP”: (B, 2, N_q, N_h, N_l, N_p)

  • attention_weights (nl.ndarray) – Attention weights in HBM, shape (B, N_q, N_h, N_l, N_p)

  • value_layout (str) – Layout of value tensor, either “BLNC” or “BNLC”. Default: “BLNC”

  • sampling_locations_layout (str) – Layout of sampling_locations, either “BQHLP2” or “B2QHLP”. Default: “BQHLP2”

  • align_corners (bool) – If True, coordinates map [0,1] to [0, H-1]. If False, map to [-0.5, H-0.5]. Default: False

  • padding_mode (str) – Padding mode for out-of-bounds coordinates, either “zeros” or “border”. Default: “zeros”

Returns:

Attention output in HBM, shape (B, N_q, N_h * C_h)

Return type:

nl.ndarray

Dimensions:

  • B: Batch size

  • N_q: Number of queries

  • N_h: Number of attention heads

  • C_h: Channels per head

  • N_l: Number of feature pyramid levels

  • N_p: Number of sampling points per query per head per level

  • L: Total flattened spatial dimension (sum of H_i * W_i across all levels)

  • H_i: Height of feature map at level i

This document is relevant for: Trn2, Trn3