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.NkiTensor, value: nl.NkiTensor, spatial_shapes: tuple, level_start_index: tuple, sampling_locations: nl.NkiTensor, attention_weights: nl.NkiTensor, value_layout: str = 'BLNC', sampling_locations_layout: str = 'BQHLP2', align_corners: bool = False, padding_mode: str = 'zeros', max_gather_indices_per_indirect: Optional[int] = None, max_scatter_indices_per_indirect: Optional[int] = None, gather_method: str = 'transpose', num_scatter_buffers: Optional[int] = None, compute_grad_value: bool = True, compute_grad_sampling_locations: bool = True, compute_grad_attention_weights: bool = True, iota_k_scale_in: Optional[nl.NkiTensor] = None, grad_value_buffers_in: Optional[nl.NkiTensor] = None, prev_reduced_grad_value: Optional[nl.NkiTensor] = None, dump_iota_k_scale: bool = False, dump_grad_value_buffers: bool = False, dump_reduced_grad_value: bool = False) → Tuple[nl.NkiTensor, ...]#

Multi-scale deformable attention backward pass kernel.

Parameters:

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

  • value (nl.NkiTensor) – 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.NkiTensor) – 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.NkiTensor) – 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”

  • max_gather_indices_per_indirect (Optional[int]) – Cap on gather indices (P_MAX * M_batch) per batched gather. Setting this (not None) enables the batched indirect gather path (M_batch = cap // P_MAX); None disables it (each (head, level) gathered separately). Default: None

  • max_scatter_indices_per_indirect (Optional[int]) – Controls the scatter-add indirection budget. If None, scatters over all P_MAX query partitions with no column batching. If set and <= P_MAX, enables the “fewer indirections” path: at most cap query partitions are scattered per dma_compute. Default: None

  • gather_method (str) – How values are gathered: “transpose” (dma_transpose to [channel, query] then nc_transpose to [query, channel]) or “copy” (indirect dma_copy to [query, channel]). Scatter always uses dma_compute. Default: “transpose”

  • num_scatter_buffers (Optional[int]) – Number of K-replicated scatter buffer sets held in HBM. Defaults to 1 when None.

  • compute_grad_value (bool) – Compute grad_value. At least one of the three compute_grad_* flags must be True. Default: True.

  • compute_grad_sampling_locations (bool) – Compute grad_sampling_locations. Default: True.

  • compute_grad_attention_weights (bool) – Compute grad_attention_weights. Default: True.

  • iota_k_scale_in (Optional[nl.NkiTensor]) – Precomputed K-scaling index tensor in HBM, shape (P_MAX, k_scale_cols), uint32. If given, the kernel copies it in instead of building it. Default: None

  • grad_value_buffers_in (Optional[nl.NkiTensor]) – Persistent K replicated scatter buffers in HBM, shape (lnc*K, L, N_h, C_h) (BLNC) / (lnc*K, N_h, L, C_h) (BNLC). If given, the kernel accumulates into it and skips zeroing; if None, the buffers are allocated and zeroed. Default: None

  • prev_reduced_grad_value (Optional[nl.NkiTensor]) – Reduced cumulative grad_value through the previous call, shape/layout of a single value map (L, N_h, C_h) / (N_h, L, C_h). When given, grad_value = reduce(K-buffers) - prev_reduced_grad_value, recovering this call’s own contribution. Default: None (grad_value = reduce(K-buffers)).

  • dump_iota_k_scale (bool) – If True, also return the built/ingested iota_k_scale in HBM so a later call can pass it as iota_k_scale_in. Default: False

  • dump_grad_value_buffers (bool) – If True, also return the raw K replicated scatter buffers in HBM so a later call can pass them as grad_value_buffers_in. Default: False

  • dump_reduced_grad_value (bool) – If True, also return reduce(K-buffers) (cumulative through this call) in HBM so a later call can pass it as prev_reduced_grad_value. Default: False

Returns:

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

Returns:

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

Returns:

Gradient w.r.t. attention_weights in HBM, shape (B, N_q, N_h, N_l, N_p) Followed, in this fixed order, by whichever of these are enabled: iota_k_scale (dump_iota_k_scale), grad_value_buffers (dump_grad_value_buffers), reduced_grad_value (dump_reduced_grad_value).

Return type:

Tuple[nl.NkiTensor, ...]

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