Skip to content

serialized_attention

Serialized attention variants used by Point Transformer V3 and descendants.

Classes:

Functions:

  • split_batch –

    Split batches into multiple sub-batches of a given size.

  • patch_layout –

    Pad every sample to whole patches of patch_size points and derive the patch offsets.

PatchLayout [source]

Bases: TypedDict

Padding and patch offsets of one point set for serialized attention.

Every block of a stage attends over the same points, so the encoder builds this once per stage with patch_layout and hands it to its blocks instead of recomputing it in every attention call.

Attributes:

  • padded_indices (Tensor) –

    Row of every padded slot, shape \((N_p,)\).

  • unpadded_indices (Tensor) –

    Padded slot of every row, shape \((N,)\).

  • padded_batch (Tensor) –

    Batch index of every padded slot, shape \((N_p,)\).

  • cu_seqlens (Tensor) –

    Cumulative patch lengths as int32, shape \((P + 1,)\), the layout flash attention takes.

RelativePositionalEncoding [source] 

RelativePositionalEncoding(patch_size: int, num_heads: int)

Bases: Module

Table-based relative position bias added to the attention logits of a serialized patch.

Each axis of the clamped relative grid coordinates indexes its own slice of a shared table, and the three per-head biases are summed. Clamping leaves \(\text{rpe\_num}^3\) distinct offsets, so the three rows are summed once per combination into a lookup table and the bias is a single gather. Summing them per point pair instead would hold \((\text{pairs}, 3)\) offsets and \((\text{pairs}, 3, H)\) gathered rows, which on a large cloud is the peak allocation of the whole model.

Parameters:

  • patch_size (int) –

    Number of points attending to each other, setting the clamping boundary.

  • num_heads (int) –

    Number of attention heads.

Shape
  • Input: \((P, K, 3)\) integer grid coordinates, \(K\) points per patch.
  • Output: \((P, H, K, K)\) bias, aligned with the attention logits.
Example 
import torch

from torch_pointcloud.layers import RelativePositionalEncoding

rpe = RelativePositionalEncoding(patch_size=8, num_heads=4)
bias = rpe(torch.randint(0, 16, (2, 8, 3)))
bias.shape  # (2, 4, 8, 8)

SerializedAttention [source] 

SerializedAttention(
    channels: int,
    num_heads: int,
    patch_size: int,
    qkv_bias: bool = True,
    qk_scale: Optional[float] = None,
    attn_drop: float = 0.0,
    proj_drop: float = 0.0,
    use_flash_attn: bool = True,
    upcast_attn: bool = True,
    upcast_softmax: bool = True,
)

Bases: Module

Vanilla serialized attention from Point Transformer V3.

No positional information is added inside attention itself: relative structure comes from the conditional position embedding (CPE) applied around each block.

Note

The input must be batch-blocked: each sample's points contiguous (a sorted batch), and any serialized_order must keep samples contiguous (the serialization codes embed the batch index). Otherwise patches silently mix points from different samples and attention leaks across the batch.

Methods:

  • patch_layout –

    The PatchLayout this layer needs for batch; build it once per stage and pass it to forward.

patch_layout [source] 

patch_layout(batch: Tensor) -> PatchLayout

The PatchLayout this layer needs for batch; build it once per stage and pass it to forward.

SerializedAttentionRPE [source] 

SerializedAttentionRPE(
    channels: int,
    num_heads: int,
    patch_size: int,
    qkv_bias: bool = True,
    qk_scale: Optional[float] = None,
    attn_drop: float = 0.0,
    proj_drop: float = 0.0,
    upcast_attn: bool = True,
    upcast_softmax: bool = True,
)

Bases: Module

Serialized attention with the relative position bias from PT-V3.

Adds a learned per-head bias indexed by the integer voxel-grid offset between query and key inside each patch. Flash Attention does not support arbitrary attention biases, so this variant always uses the manual softmax path.

Note

The input must be batch-blocked: each sample's points contiguous (a sorted batch), and any serialized_order must keep samples contiguous (the serialization codes embed the batch index). Otherwise patches silently mix points from different samples and attention leaks across the batch.

Methods:

  • patch_layout –

    The PatchLayout this layer needs for batch; build it once per stage and pass it to forward.

patch_layout [source] 

patch_layout(batch: Tensor) -> PatchLayout

The PatchLayout this layer needs for batch; build it once per stage and pass it to forward.

SerializedAttentionRoPE [source] 

SerializedAttentionRoPE(
    channels: int,
    num_heads: int,
    patch_size: int,
    qkv_bias: bool = True,
    qk_scale: Optional[float] = None,
    attn_drop: float = 0.0,
    proj_drop: float = 0.0,
    use_flash_attn: bool = True,
    upcast_attn: bool = True,
    upcast_softmax: bool = True,
    rope_base: float = 10.0,
)

Bases: Module

Serialized attention with 3D rotary position embedding from Utonia.

Rotates \(Q\), \(K\) via Point3DRoPE using the real-valued metric position of each token. Flash Attention is supported and uses bfloat16 (matching upstream Utonia's reference implementation).

Note

The input must be batch-blocked: each sample's points contiguous (a sorted batch), and any serialized_order must keep samples contiguous (the serialization codes embed the batch index). Otherwise patches silently mix points from different samples and attention leaks across the batch.

Methods:

  • patch_layout –

    The PatchLayout this layer needs for batch; build it once per stage and pass it to forward.

patch_layout [source] 

patch_layout(batch: Tensor) -> PatchLayout

The PatchLayout this layer needs for batch; build it once per stage and pass it to forward.

split_batch [source] 

split_batch(batch: Tensor, max_size: int) -> Tensor

Split batches into multiple sub-batches of a given size.

Note

The batch is only splitted if it is larger than the given size. If not, the batch is returned as is.

Note

If you want to split batches smaller than the given size, you can use the divisible_pad function before splitting the batch.

Parameters:

  • batch (Tensor) –

    The batch indices of the points.

  • max_size (int) –

    The maximum size of the sub-batches.

Returns:

  • Tensor –

    The sub-batch indices.

Examples:

>>> import torch
>>> batch = torch.tensor([0, 0, 0, 1, 1, 1, 1, 2, 2, 3])
>>> split_batch(batch, max_size=2)
tensor([0, 0, 1, 2, 2, 3, 3, 4, 4, 5])

patch_layout [source] 

patch_layout(
    batch: Tensor, patch_size: int, mode: PadMode
) -> PatchLayout

Pad every sample to whole patches of patch_size points and derive the patch offsets.

Flash attention handles a sample smaller than the patch as a short variable-length patch, so it only pads samples at or above the patch size (mode="above"); the dense reshape needs every sample padded to a full patch (mode="all"). Padding, never shrinking, keeps one sample's output independent of its co-batched neighbors.

Parameters:

  • batch (Tensor) –

    Batch index of every point, sorted, shape \((N,)\).

  • patch_size (int) –

    Number of points per attention patch.

  • mode (PadMode) –

    Which samples to pad, see divisible_pad.

Returns: