# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
# Copyright 2026 Tensor Auto Inc. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
PaliGemma with Expert Module.
This module implements the PaliGemma model with an additional expert module,
specifically designed for the Pi05 policy. It combines a pre-trained PaliGemma
Vision-Language Model (VLM) with a Gemma-based expert model to handle
action generation and conditioning.
"""
import logging
import torch
from einops import rearrange
from torch import nn
from transformers import (
AutoConfig,
Cache,
GemmaForCausalLM,
PaliGemmaForConditionalGeneration,
PretrainedConfig,
PreTrainedModel,
)
from transformers.models.auto import CONFIG_MAPPING
from transformers.models.gemma import modeling_gemma
from opentau.policies import flash_attn_cuda
def _preferred_dtype():
return torch.float32 if torch.onnx.is_in_onnx_export() else torch.bfloat16
[docs]
def apply_rope(x: torch.Tensor, positions: torch.Tensor, max_wavelength: int = 10_000) -> torch.Tensor:
"""Applies (multimodal) RoPE positions to the input tensor.
Two position layouts are supported, selected by ``positions.ndim``:
* **1-D RoPE** — ``positions`` of shape ``[B, L]``: a single scalar
position per token (the historical behaviour). Every rotary frequency
band is driven by that one position.
* **Interleaved MRoPE** — ``positions`` of shape ``[A, B, L]`` (here
``A == 3`` for the temporal / height / width axes): frequency band ``i``
is driven by axis ``i % A``, so the three axes are *interleaved* across
the full frequency spectrum rather than each owning a contiguous block
(the Qwen3-VL "interleaved MRoPE" scheme). When all ``A`` axes carry the
same positions this is bit-identical to the 1-D path, so text-style
tokens (``t == h == w``) rotate exactly as plain RoPE.
Args:
x: Input tensor of shape [B, L, H, D].
positions: Position tensor of shape ``[B, L]`` (1-D RoPE) or
``[A, B, L]`` (interleaved MRoPE).
max_wavelength: Maximum wavelength for RoPE. Defaults to 10_000.
Returns:
Tensor: The input tensor with RoPE applied, of shape [B, L, H, D].
"""
d_half = x.shape[-1] // 2
device = x.device
dtype = x.dtype
x = x.to(torch.float32)
freq_exponents = (2.0 / x.shape[-1]) * torch.arange(d_half, dtype=torch.float32, device=device)
timescale = max_wavelength**freq_exponents
if positions.ndim == 2:
# 1-D RoPE: [B, L] -> [B, L, d_half]
radians = positions[..., None].to(torch.float32) / timescale[None, None, :].to(torch.float32)
else:
# Interleaved MRoPE: positions [A, B, L]. Assign frequency band i to
# axis i % A so every axis spans the full spectrum, then select the
# per-band position -> [B, L, d_half].
n_axes = positions.shape[0]
band_axis = torch.arange(d_half, device=device) % n_axes # [d_half]
pos_per_band = positions.to(torch.float32)[band_axis] # [d_half, B, L]
pos_per_band = rearrange(pos_per_band, "f b l -> b l f") # [B, L, d_half]
radians = pos_per_band / timescale[None, None, :].to(torch.float32)
radians = radians[..., None, :]
sin = torch.sin(radians) # .to(dtype=dtype)
cos = torch.cos(radians) # .to(dtype=dtype)
x1, x2 = x.split(d_half, dim=-1)
res = torch.empty_like(x)
res[..., :d_half] = x1 * cos - x2 * sin
res[..., d_half:] = x2 * cos + x1 * sin
return res.to(dtype)
[docs]
class PaliGemmaWithExpertConfig(PretrainedConfig):
"""Configuration class for PaliGemmaWithExpertModel."""
model_type = "PaliGemmaWithExpertModel"
sub_configs = {"paligemma_config": AutoConfig, "gemma_expert_config": AutoConfig}
[docs]
def __init__(
self,
paligemma_config: dict | None = None,
gemma_expert_config: dict | None = None,
freeze_vision_encoder: bool = True,
train_expert_only: bool = True,
attention_implementation: str = "eager",
discrete_action_vocab_size: int | None = None,
dropout: float = 0.1,
gradient_checkpointing: bool = False,
**kwargs,
):
"""Initializes the configuration.
Args:
paligemma_config: Configuration dictionary for the PaliGemma model.
gemma_expert_config: Configuration dictionary for the Gemma expert model.
freeze_vision_encoder: Whether to freeze the vision encoder. Defaults to True.
train_expert_only: Whether to train only the expert model. Defaults to True.
attention_implementation: Attention implementation to use ("eager", "sdpa",
or "fa2"). Defaults to "eager".
discrete_action_vocab_size: Vocabulary size for discrete actions.
dropout: Dropout probability. Defaults to 0.1.
gradient_checkpointing: Wrap each decoder-layer body in
``torch.utils.checkpoint.checkpoint`` during training. Trades
roughly one extra forward pass per step (~25-33% compute)
for ~30-40 GB of activation memory per rank, which enables
larger per-rank batch sizes and amortizes per-step fixed cost.
Only safe under plain DDP (MULTI_GPU), single-process (NO),
or DeepSpeed ZeRO-1/2 — see the train.py guard. Defaults to False.
**kwargs: Additional keyword arguments passed to PretrainedConfig.
"""
self.freeze_vision_encoder = freeze_vision_encoder
self.train_expert_only = train_expert_only
self.attention_implementation = attention_implementation
self.discrete_action_vocab_size = discrete_action_vocab_size
self.dropout = dropout
self.gradient_checkpointing = gradient_checkpointing
if paligemma_config is None:
# Default config from Pi0
self.paligemma_config = CONFIG_MAPPING["paligemma"](
transformers_version="4.48.1",
_vocab_size=257152,
bos_token_id=2,
eos_token_id=1,
hidden_size=2048,
image_token_index=257152,
model_type="paligemma",
pad_token_id=0,
projection_dim=2048,
text_config={
"hidden_activation": "gelu_pytorch_tanh",
"hidden_size": 2048,
"intermediate_size": 16384,
"model_type": "gemma",
"num_attention_heads": 8,
"num_hidden_layers": 18,
"num_image_tokens": 256,
"num_key_value_heads": 1,
"torch_dtype": "float32",
"vocab_size": 257152,
"use_adarms": False,
"adarms_cond_dim": None,
},
vision_config={
"hidden_size": 1152,
"intermediate_size": 4304,
"model_type": "siglip_vision_model",
"num_attention_heads": 16,
"num_hidden_layers": 27,
"num_image_tokens": 256,
"patch_size": 14,
"projection_dim": 2048,
"projector_hidden_act": "gelu_fast",
"torch_dtype": "float32",
"vision_use_head": False,
},
)
elif isinstance(self.paligemma_config, dict):
# Override Pi0 default config for PaliGemma
if "model_type" not in gemma_expert_config:
paligemma_config["model_type"] = "paligemma"
cfg_cls = CONFIG_MAPPING[paligemma_config["model_type"]]
self.paligemma_config = cfg_cls(**paligemma_config)
if gemma_expert_config is None:
# Default config from Pi0
self.gemma_expert_config = CONFIG_MAPPING["gemma"](
attention_bias=False,
attention_dropout=0.0,
bos_token_id=2,
eos_token_id=1,
head_dim=256,
hidden_act="gelu_pytorch_tanh",
hidden_activation="gelu_pytorch_tanh",
hidden_size=1024,
initializer_range=0.02,
intermediate_size=4096,
max_position_embeddings=8192,
model_type="gemma",
num_attention_heads=8,
num_hidden_layers=18,
num_key_value_heads=1,
pad_token_id=0,
rms_norm_eps=1e-06,
rope_theta=10000.0,
torch_dtype="float32",
use_adarms=True,
adarms_cond_dim=1024,
transformers_version="4.48.1",
use_cache=True,
vocab_size=257152,
)
elif isinstance(self.gemma_expert_config, dict):
# Override Pi0 default config for Gemma Expert
if "model_type" not in gemma_expert_config:
gemma_expert_config["model_type"] = "gemma"
cfg_cls = CONFIG_MAPPING[paligemma_config["model_type"]]
self.gemma_expert_config = cfg_cls(**gemma_expert_config)
super().__init__(**kwargs)
def __post_init__(self):
"""Validates configuration parameters."""
super().__post_init__()
if self.train_expert_only and not self.freeze_vision_encoder:
raise ValueError(
"You set `freeze_vision_encoder=False` and `train_expert_only=True` which are not compatible."
)
if self.attention_implementation not in ["eager", "sdpa", "fa2", "flash_cuda"]:
raise ValueError(
f"Wrong value provided for `attention_implementation` ({self.attention_implementation}). "
"Expected 'eager', 'sdpa', 'fa2', or 'flash_cuda'."
)
if self.attention_implementation == "fa2":
# "fa2" has been accepted by the validator historically but never
# implemented in PaliGemmaWithExpertModel. Fall back to eager so
# existing configs keep running.
logging.warning(
"attention_implementation='fa2' is not implemented; falling back to 'eager'. "
"Consider switching to 'sdpa' for ~10-15% better throughput."
)
[docs]
class PaliGemmaWithExpertModel(PreTrainedModel):
"""PaliGemma model with an additional expert module for action generation."""
config_class = PaliGemmaWithExpertConfig
[docs]
def __init__(self, config: PaliGemmaWithExpertConfig):
"""Initializes the PaliGemmaWithExpertModel.
Args:
config: Configuration object for the model.
"""
super().__init__(config=config)
self.config = config
self.paligemma = PaliGemmaForConditionalGeneration(config=config.paligemma_config)
self.gemma_expert = GemmaForCausalLM(config=config.gemma_expert_config)
# Remove unused embed_tokens
self.gemma_expert.model.embed_tokens = None
# Remove unused lm_head. The action expert is used as an encoder: its
# output is routed through action_out_proj for continuous flow-matching
# actions and through da_head for discrete-action CE. It never flows
# through a token-vocabulary head. Leaving lm_head in place registers a
# ~263M-param weight that receives no gradient and forces
# DistributedDataParallelKwargs(find_unused_parameters=True), which costs
# a per-step DDP graph walk (~10-15% of step time).
self.gemma_expert.lm_head = None
# Learned embedding layer for discrete actions
# Embedding dimension matches expert model hidden size
self.discrete_action_embedding = nn.Embedding(
num_embeddings=config.discrete_action_vocab_size,
embedding_dim=config.paligemma_config.text_config.hidden_size,
padding_idx=0, # 0 is used for padding in pad_fast_tokens
)
# discrete action head that maps to action vocab size and not language vocab size
self.da_head = nn.Linear(
in_features=config.paligemma_config.text_config.hidden_size,
out_features=config.discrete_action_vocab_size,
)
self.dropout = nn.Dropout(config.dropout)
if not torch.compiler.is_compiling(): # Only cast to bfloat16 if not compiling
self.to_bfloat16_like_physical_intelligence()
self.set_requires_grad()
[docs]
def set_requires_grad(self) -> None:
"""Sets the requires_grad attribute for model parameters based on configuration."""
if self.config.freeze_vision_encoder:
self.paligemma.vision_tower.eval()
for params in self.paligemma.vision_tower.parameters():
params.requires_grad = False
if self.config.train_expert_only:
self.paligemma.eval()
for params in self.paligemma.parameters():
params.requires_grad = False
for param in self.da_head.parameters():
param.requires_grad = False
for param in self.discrete_action_embedding.parameters():
param.requires_grad = False
[docs]
def train(self, mode: bool = True) -> None:
"""Sets the module in training mode.
Args:
mode: whether to set training mode (True) or evaluation mode (False). Defaults to True.
"""
super().train(mode)
if self.config.freeze_vision_encoder:
self.paligemma.vision_tower.eval()
if self.config.train_expert_only:
self.paligemma.eval()
[docs]
def to_bfloat16_like_physical_intelligence(self) -> None:
"""Casts specific model components to bfloat16 dtype."""
self.paligemma = self.paligemma.to(dtype=torch.bfloat16)
params_to_change_dtype = [
"language_model.model.layers",
"gemma_expert.model.layers",
"vision_tower",
"multi_modal",
]
for name, param in self.named_parameters():
if any(selector in name for selector in params_to_change_dtype):
param.data = param.data.to(dtype=torch.bfloat16)
[docs]
def embed_image(self, image: torch.Tensor) -> torch.Tensor:
"""Computes image embeddings.
Args:
image: Input image tensor.
Returns:
torch.Tensor: Image embeddings.
"""
# Handle different transformers versions
if hasattr(self.paligemma, "get_image_features"):
return self.paligemma.get_image_features(image)
else:
return self.paligemma.model.get_image_features(image)
[docs]
def embed_language_tokens(self, tokens: torch.Tensor) -> torch.Tensor:
"""Embeds language tokens.
Args:
tokens: Input token indices.
Returns:
torch.Tensor: Token embeddings.
"""
return self.paligemma.language_model.embed_tokens(tokens)
[docs]
def embed_discrete_actions(self, actions: torch.Tensor) -> torch.Tensor:
"""Embeds discrete action tokens.
Args:
actions: Input discrete action indices.
Returns:
torch.Tensor: Action embeddings.
"""
# Ensure actions are long integers for embedding lookup
if actions.dtype != torch.long:
actions = actions.long()
# Apply embedding layer
embedded = self.discrete_action_embedding(actions)
return embedded
# TODO: break down this huge forward into modules or functions
[docs]
def forward(
self,
attention_mask: torch.Tensor | None = None,
position_ids: torch.LongTensor | None = None,
past_key_values: list[torch.FloatTensor] | Cache | None = None,
inputs_embeds: list[torch.FloatTensor] = None,
n_cross_att_tokens: int | None = None,
use_cache: bool | None = None,
fill_kv_cache: bool | None = None,
adarms_cond: list[torch.Tensor] | None = None,
attention_block_ids: tuple[torch.Tensor, ...] | None = None,
) -> tuple[list[torch.FloatTensor | None], list[torch.FloatTensor] | Cache | None]:
"""Forward pass of the model.
Args:
attention_mask: Dense bool attention mask (B, Sq, Sk), True=attend.
Used by the eager/sdpa backends. ``None`` when the ``flash_cuda``
backend is active (the mask is reconstructed in-kernel instead).
attention_block_ids: Compact ``(q_blk, k_blk, q_valid, k_valid)``
block-id representation for the ``flash_cuda`` backend (see
``flash_attn_cuda.make_att_block_ids``). ``None`` for eager/sdpa.
position_ids: Position IDs tensor.
past_key_values: Past key values for caching.
inputs_embeds: List of input embeddings for the different model parts.
n_cross_att_tokens: Number of cross-attention tokens.
use_cache: Whether to use KV cache.
fill_kv_cache: Whether to fill the KV cache.
adarms_cond: List of AdaRMS conditioning tensors.
Returns:
tuple: A tuple containing:
- outputs_embeds: List of output embeddings.
- past_key_values: Updated past key values.
Raises:
ValueError: If `n_cross_att_tokens` is not provided when `fill_kv_cache` is True.
"""
if adarms_cond is None:
adarms_cond = [None, None]
models = [self.paligemma.language_model, self.gemma_expert.model]
for hidden_states in inputs_embeds:
# TODO this is very inefficient
# dtype is always the same, batch size too (if > 1 len)
# device could be trickier in multi gpu edge cases but that's it
if hidden_states is None:
continue
batch_size = hidden_states.shape[0]
# RMSNorm
num_layers = self.paligemma.config.text_config.num_hidden_layers
head_dim = self.paligemma.config.text_config.head_dim
# If gradient checkpointing will be writing to past_key_values, make
# sure the dict exists before we enter the loop. The original code
# lazily created it on the first layer via
# ``if past_key_values is None: past_key_values = {}``; doing it here
# hoists that out so ``_run_layer`` always receives a non-None dict
# when fill_kv_cache=True, which is important for checkpoint recompute
# to be idempotent.
if fill_kv_cache and past_key_values is None:
past_key_values = {}
use_ckpt = self.config.gradient_checkpointing and self.training
for layer_idx in range(num_layers):
if use_ckpt:
# use_reentrant=False is the modern, DDP-safe path; it
# preserves RNG state across recompute so dropout is
# deterministic, and participates cleanly in autograd's
# saved_tensors_hooks.
inputs_embeds = torch.utils.checkpoint.checkpoint(
self._run_layer,
layer_idx,
inputs_embeds,
attention_mask,
position_ids,
past_key_values,
n_cross_att_tokens,
use_cache,
fill_kv_cache,
adarms_cond,
batch_size,
head_dim,
attention_block_ids,
use_reentrant=False,
)
else:
inputs_embeds = self._run_layer(
layer_idx,
inputs_embeds,
attention_mask,
position_ids,
past_key_values,
n_cross_att_tokens,
use_cache,
fill_kv_cache,
adarms_cond,
batch_size,
head_dim,
attention_block_ids,
)
# final norm
outputs_embeds = []
for i, hidden_states in enumerate(inputs_embeds):
if hidden_states is not None:
out_emb, _ = models[i].norm(hidden_states, cond=adarms_cond[i])
outputs_embeds.append(out_emb)
else:
outputs_embeds.append(None)
return outputs_embeds, past_key_values
def _run_layer(
self,
layer_idx: int,
inputs_embeds: list[torch.FloatTensor | None],
attention_mask: torch.Tensor | None,
position_ids: torch.LongTensor | None,
past_key_values: dict | None,
n_cross_att_tokens: int | None,
use_cache: bool | None,
fill_kv_cache: bool | None,
adarms_cond: list[torch.Tensor | None],
batch_size: int,
head_dim: int,
attention_block_ids: tuple[torch.Tensor, ...] | None = None,
) -> list[torch.FloatTensor | None]:
"""Run a single layer of the dual-tower decoder loop.
Extracted from ``forward()`` as a standalone method so it can be the
unit of ``torch.utils.checkpoint.checkpoint`` wrapping when
``config.gradient_checkpointing`` is enabled. Behavior is bit-identical
to the original inlined loop body; side effects on ``past_key_values``
are preserved (mutation is idempotent across checkpoint recompute
because each layer writes its own unique key).
Args:
layer_idx: Index of the current decoder layer.
inputs_embeds: Per-tower input embeddings. Entries may be None if
that tower is not participating in this forward (e.g. expert
tower during prefix-only pass).
attention_mask: Bool attention mask, True = attend, shape (B, S, S).
position_ids: Position IDs tensor passed through RoPE.
past_key_values: KV cache dict, populated if fill_kv_cache, read
if use_cache.
n_cross_att_tokens: Number of prefix tokens to cache for cross-
attention (required when fill_kv_cache).
use_cache: If True, prepend cached KV for this layer.
fill_kv_cache: If True, write this layer's K/V into the cache.
adarms_cond: Per-tower AdaRMSNorm conditioning tensors (or None).
batch_size: Cached batch size from the outer forward.
head_dim: Per-head dimension.
Returns:
list[torch.FloatTensor | None]: Per-tower output embeddings for
this layer.
"""
models = [self.paligemma.language_model, self.gemma_expert.model]
query_states = []
key_states = []
value_states = []
gates = []
for i, hidden_states in enumerate(inputs_embeds):
if hidden_states is None:
gates.append(None)
continue
layer = models[i].layers[layer_idx]
hidden_states, gate = layer.input_layernorm(hidden_states, cond=adarms_cond[i])
gates.append(gate)
input_shape = hidden_states.shape[:-1]
hidden_shape = (*input_shape, -1, layer.self_attn.head_dim)
hidden_states = hidden_states.to(dtype=_preferred_dtype())
query_state = layer.self_attn.q_proj(hidden_states).view(hidden_shape)
key_state = layer.self_attn.k_proj(hidden_states).view(hidden_shape)
value_state = layer.self_attn.v_proj(hidden_states).view(hidden_shape)
query_states.append(query_state)
key_states.append(key_state)
value_states.append(value_state)
# B,L,H,D with L sequence length, H number of heads, D head dim
# concatenate on the number of embeddings/tokens
query_states = torch.cat(query_states, dim=1)
key_states = torch.cat(key_states, dim=1)
value_states = torch.cat(value_states, dim=1)
query_states = apply_rope(query_states, position_ids)
key_states = apply_rope(key_states, position_ids)
if use_cache:
# TODO here, some optimization can be done - similar to a `StaticCache` we can declare the `max_len` before.
# so we create an empty cache, with just one cuda malloc, and if (in autoregressive case) we reach
# the max len, then we (for instance) double the cache size. This implementation already exists
# in `transformers`. (molbap)
key_states = torch.cat([past_key_values[layer_idx]["key_states"], key_states], dim=1)
value_states = torch.cat([past_key_values[layer_idx]["value_states"], value_states], dim=1)
if fill_kv_cache:
if n_cross_att_tokens is None:
raise ValueError("n_cross_att_tokens must be provided when fill_kv_cache is True")
past_key_values[layer_idx] = {
# save the first n_cross_att_tokens for action expert cross attention
"key_states": key_states[:, :n_cross_att_tokens, :, :],
"value_states": value_states[:, :n_cross_att_tokens, :, :],
}
if attention_block_ids is not None:
# flash_cuda backend: mask reconstructed in-kernel from block-ids.
att_output = self.flash_attention_forward(
attention_block_ids, batch_size, head_dim, query_states, key_states, value_states
)
else:
attention_interface = self.get_attention_interface()
att_output = attention_interface(
attention_mask, batch_size, head_dim, query_states, key_states, value_states
)
att_output = att_output.to(dtype=_preferred_dtype())
# first part of att_output is prefix (up to sequence length, [:, 0:prefix_seq_len])
outputs_embeds: list[torch.FloatTensor | None] = []
start = 0
for i, hidden_states in enumerate(inputs_embeds):
layer = models[i].layers[layer_idx]
if hidden_states is not None:
end = start + hidden_states.shape[1]
if att_output.dtype != layer.self_attn.o_proj.weight.dtype:
att_output = att_output.to(layer.self_attn.o_proj.weight.dtype)
out_emb = layer.self_attn.o_proj(att_output[:, start:end])
out_emb = self.dropout(out_emb)
# first residual
out_emb = modeling_gemma._gated_residual(hidden_states, out_emb, gates[i]) # noqa: SLF001
after_first_residual = out_emb.clone()
out_emb, gate = layer.post_attention_layernorm(out_emb, cond=adarms_cond[i])
out_emb = layer.mlp(out_emb)
out_emb = self.dropout(out_emb)
# second residual
out_emb = modeling_gemma._gated_residual(after_first_residual, out_emb, gate) # noqa: SLF001
outputs_embeds.append(out_emb)
start = end
else:
outputs_embeds.append(None)
return outputs_embeds
[docs]
def get_attention_interface(self):
"""Returns the attention implementation function based on config.
Dispatches on ``self.config.attention_implementation``:
- ``"eager"``: naive matmul-softmax-matmul in fp32 (historical
default; see ``eager_attention_forward``).
- ``"sdpa"``: ``torch.nn.functional.scaled_dot_product_attention``
which on A100 dispatches to FlashAttention-2 or the
memory-efficient backend — 2-3x faster than eager at pi05's
sequence length.
- ``"fa2"``: accepted for backward compatibility; falls back to
eager with a warning emitted at config validation time.
- ``"flash_cuda"``: custom block-causal CUDA flash kernel, dispatched
in ``_run_layer`` via ``flash_attention_forward`` whenever
``attention_block_ids`` is supplied. There is **no** silent fallback:
if this method is ever reached under ``flash_cuda`` it means a caller
failed to thread the block-ids, which is a bug, so it raises.
Returns:
callable: The attention function to use.
Raises:
RuntimeError: If ``attention_implementation == "flash_cuda"`` (the
flash path must go through ``flash_attention_forward`` with
block-ids, not this interface).
"""
impl = self.config.attention_implementation
if impl == "flash_cuda":
raise RuntimeError(
"flash_cuda must be dispatched via flash_attention_forward with "
"attention_block_ids; get_attention_interface reached instead "
"(block-ids were not threaded to this forward)."
)
if impl == "sdpa":
return self.sdpa_attention_forward
# "eager" and legacy "fa2" both land here; "fa2" already warned
# during __post_init__.
return self.eager_attention_forward
[docs]
def flash_attention_forward(
self,
attention_block_ids: tuple[torch.Tensor, ...],
batch_size: int,
head_dim: int,
query_states: torch.Tensor,
key_states: torch.Tensor,
value_states: torch.Tensor,
) -> torch.Tensor:
"""Custom block-causal CUDA flash attention forward.
Unlike the eager/sdpa interfaces this consumes the compact block-id
representation (``q_blk, k_blk, q_valid, k_valid``) and reconstructs the
block-causal mask inside the kernel, so the dense (B, Sq, Sk) mask is
never materialized. GQA/MQA is handled natively (no K/V expansion), and
the kernel runs both matmuls on Tensor Cores (fp16/bf16) with fp32
accumulation, matching ``eager``/``sdpa`` outputs within fp/bf16 noise.
Args:
attention_block_ids: ``(q_blk, k_blk, q_valid, k_valid)`` from
``flash_attn_cuda.make_att_block_ids``.
batch_size: Batch size.
head_dim: Per-head dimension.
query_states: ``(B, Sq, num_attention_heads, head_dim)``.
key_states: ``(B, Sk, num_key_value_heads, head_dim)``.
value_states: ``(B, Sk, num_key_value_heads, head_dim)``.
Returns:
torch.Tensor: ``(B, Sq, num_attention_heads * head_dim)`` output,
matching the eager/sdpa interfaces.
"""
q_blk, k_blk, q_valid, k_valid = attention_block_ids
att_output = flash_attn_cuda.flash_attn_blockmask(
query_states, key_states, value_states, q_blk, k_blk, q_valid, k_valid, head_dim**-0.5
)
# (B, Sq, H, head_dim) -> (B, Sq, H * head_dim) to match eager/sdpa.
return att_output.reshape(batch_size, att_output.shape[1], -1)
[docs]
def eager_attention_forward(
self,
attention_mask: torch.Tensor,
batch_size: int,
head_dim: int,
query_states: torch.Tensor,
key_states: torch.Tensor,
value_states: torch.Tensor,
) -> torch.Tensor:
"""Eager attention forward pass using standard matrix multiplications.
Args:
attention_mask: Attention mask tensor.
batch_size: Batch size.
head_dim: Head dimension.
query_states: Query states tensor.
key_states: Key states tensor.
value_states: Value states tensor.
Returns:
torch.Tensor: Attention output.
"""
num_att_heads = self.config.paligemma_config.text_config.num_attention_heads
num_key_value_heads = self.config.paligemma_config.text_config.num_key_value_heads
num_key_value_groups = num_att_heads // num_key_value_heads
# query_states: batch_size, sequence_length, num_att_head, head_dim
# key_states: batch_size, sequence_length, num_key_value_head, head_dim
# value_states: batch_size, sequence_length, num_key_value_head, head_dim
sequence_length = key_states.shape[1]
key_states = key_states[:, :, :, None, :].expand(
batch_size, sequence_length, num_key_value_heads, num_key_value_groups, head_dim
)
key_states = key_states.reshape(
batch_size, sequence_length, num_key_value_heads * num_key_value_groups, head_dim
)
value_states = value_states[:, :, :, None, :].expand(
batch_size, sequence_length, num_key_value_heads, num_key_value_groups, head_dim
)
value_states = value_states.reshape(
batch_size, sequence_length, num_key_value_heads * num_key_value_groups, head_dim
)
# Attention here is upcasted to float32 to match the original eager implementation.
query_states = query_states.to(dtype=torch.float32)
key_states = key_states.to(dtype=torch.float32)
query_states = query_states.transpose(1, 2)
key_states = key_states.transpose(1, 2)
att_weights = torch.matmul(query_states, key_states.transpose(2, 3))
att_weights *= head_dim**-0.5
big_neg = -2.3819763e38 # See gemma/modules.py
masked_att_weights = torch.where(attention_mask[:, None, :, :], att_weights, big_neg)
probs = nn.functional.softmax(masked_att_weights, dim=-1)
probs = probs.to(dtype=value_states.dtype)
# probs: batch_size, num_key_value_head, num_att_head, sequence_length, sequence_length
# value_states: batch_size, sequence_length, num_att_heads, head_dim
att_output = torch.matmul(probs, value_states.permute(0, 2, 1, 3))
att_output = att_output.permute(0, 2, 1, 3)
# we use -1 because sequence length can change
att_output = att_output.reshape(batch_size, -1, num_key_value_heads * num_key_value_groups * head_dim)
return att_output
[docs]
def sdpa_attention_forward(
self,
attention_mask: torch.Tensor,
batch_size: int,
head_dim: int,
query_states: torch.Tensor,
key_states: torch.Tensor,
value_states: torch.Tensor,
) -> torch.Tensor:
"""SDPA attention forward pass using ``F.scaled_dot_product_attention``.
Produces the same output shape and semantic as ``eager_attention_forward``
but delegates the scores-softmax-matmul chain to PyTorch's fused SDPA
kernel. On A100 + bf16 PyTorch typically dispatches to FlashAttention-2,
which keeps the S×S attention matrix in on-chip SRAM and runs the
matmul in bf16 with fp32 accumulation in the softmax. There is a
deliberate numerical difference vs. the eager kernel: we do **not**
upcast Q/K to float32 before the matmul, because modern attention
kernels do fp32 accumulation internally — cleaner and faster. Training
dynamics are equivalent within bf16 reassociation noise.
Args:
attention_mask: Boolean mask of shape (B, S, S); ``True`` = attend.
batch_size: Batch size.
head_dim: Per-head dimension.
query_states: Query states of shape (B, S, num_attention_heads, head_dim).
key_states: Key states of shape (B, S, num_key_value_heads, head_dim).
value_states: Value states of shape (B, S, num_key_value_heads, head_dim).
Returns:
torch.Tensor: Attention output of shape
(B, S, num_attention_heads * head_dim).
"""
num_att_heads = self.config.paligemma_config.text_config.num_attention_heads
num_key_value_heads = self.config.paligemma_config.text_config.num_key_value_heads
num_key_value_groups = num_att_heads // num_key_value_heads
sequence_length = key_states.shape[1]
# GQA expansion mirroring eager_attention_forward. Always-correct
# across PyTorch versions; for PyTorch 2.5+ we could alternatively
# pass the un-expanded K/V with enable_gqa=True to SDPA, but the
# explicit expand+reshape is a memory-view and adds no real cost.
key_states = key_states[:, :, :, None, :].expand(
batch_size, sequence_length, num_key_value_heads, num_key_value_groups, head_dim
)
key_states = key_states.reshape(
batch_size, sequence_length, num_key_value_heads * num_key_value_groups, head_dim
)
value_states = value_states[:, :, :, None, :].expand(
batch_size, sequence_length, num_key_value_heads, num_key_value_groups, head_dim
)
value_states = value_states.reshape(
batch_size, sequence_length, num_key_value_heads * num_key_value_groups, head_dim
)
# SDPA expects (B, H, S, D_h) for Q, K, V.
query_states = query_states.transpose(1, 2)
key_states = key_states.transpose(1, 2)
value_states = value_states.transpose(1, 2)
# PyTorch's SDPA accepts a bool ``attn_mask`` where True = attend,
# which matches our convention from make_att_2d_masks. Broadcast a
# single-head dim so the same mask applies to every attention head.
attn_mask = attention_mask[:, None, :, :]
# is_causal must be False: our mask already encodes both the intra-
# prefix bidirectional pattern and any causal tail. SDPA's
# ``is_causal=True`` shortcut would double-apply and be wrong.
att_output = nn.functional.scaled_dot_product_attention(
query_states,
key_states,
value_states,
attn_mask=attn_mask,
is_causal=False,
dropout_p=0.0,
)
# att_output: (B, H, S, D_h) → (B, S, H * D_h) to match eager output.
att_output = att_output.permute(0, 2, 1, 3)
att_output = att_output.reshape(batch_size, -1, num_key_value_heads * num_key_value_groups * head_dim)
return att_output