Overview

PyTorch 2.7, released on April 23, 2025, adds Blackwell GPU support, Mega Cache for torch.compile, and FlexAttention.

Main Features

Blackwell GPU support and Mega Cache

PyTorch 2.7 natively supports NVIDIA Blackwell GPUs (B100/B200). Mega Cache speeds up torch.compile by caching compiled graphs across runs.

python 
import torch

# Mega Cache: persistent compilation
torch._dynamo.config.cache_size_limit = 256

@torch.compile
def forward(x):
    return x @ x.T + torch.relu(x)

x = torch.randn(1024, 1024, device='cuda')
out = forward(x)  # compiles and caches
print(out.shape)  # torch.Size([1024, 1024])

FlexAttention

FlexAttention lets you define custom attention mechanisms with arbitrary masks while retaining Flash Attention kernel optimizations.

python 
from torch.nn.attention.flex_attention import (
    flex_attention, create_block_mask
)

# Custom causal mask
def causal_mask(b, h, q_idx, kv_idx):
    return q_idx >= kv_idx

mask = create_block_mask(causal_mask, B=1, H=1, Q_LEN=512, KV_LEN=512)
# flex_attention(query, key, value, block_mask=mask)

Sources