This document is relevant for: Trn2, Trn3

MS Deformable Attention Backward Kernel API Reference#

Multi-scale deformable attention backward pass kernel.

Computes gradients with respect to value, sampling_locations, and attention_weights given the downstream gradient.

Background#

The ms_deformable_attention_bwd kernel computes the backward pass for multi-scale deformable attention, producing gradients with respect to value, sampling_locations, and attention_weights from the downstream gradient using bilinear-interpolation derivatives. It 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_bwd.py

ms_deformable_attention_bwd#

nkilib.experimental.deformable_attention.ms_deformable_attention_bwd(grad_output: nl.ndarray, 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') → Tuple[nl.ndarray, nl.ndarray, nl.ndarray]#

Multi-scale deformable attention backward pass kernel.

Parameters:

  • grad_output (nl.ndarray) – Gradient from downstream in HBM, shape (B, N_q, N_h * C_h)

  • 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:

Gradient w.r.t. value in HBM, same shape and layout as input value

Return type:

nl.ndarray

Returns:

Gradient w.r.t. sampling_locations in HBM, same shape and layout as input

Return type:

nl.ndarray

Returns:

Gradient w.r.t. attention_weights in HBM, shape (B, N_q, N_h, N_l, N_p)

Return type:

nl.ndarray

Notes:

  • Computes actual gradients using bilinear interpolation derivatives

  • Supports both BLNC and BNLC value layouts

  • Supports both BQHLP2 and B2QHLP sampling_locations layouts

  • Padding modes: “zeros” (OOB returns 0) and “border” (clamps to edge)

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