Vue d'ensemble

PyTorch 2.7, publie le 23 avril 2025, ajoute le support des GPU Blackwell, le Mega Cache pour torch.compile et FlexAttention.

Fonctionnalites principales

Support GPU Blackwell et Mega Cache

PyTorch 2.7 supporte nativement les GPU NVIDIA Blackwell (B100/B200). Le Mega Cache accelere torch.compile en mettant en cache les graphes compiles entre les executions.

python 
import torch

# Mega Cache : compilation persistante
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)  # compile et cache
print(out.shape)  # torch.Size([1024, 1024])

FlexAttention

FlexAttention permet de definir des mecanismes d'attention personnalises avec des masques arbitraires, tout en conservant les optimisations du noyau Flash Attention.

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

# Masque causal personnalise
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