importtorchfromtorch.nn.attention.flex_attentionimportflex_attention,create_block_maskdefcausal_mask(b,h,q_idx,kv_idx):returnq_idx>=kv_idx# Setupdevice=torch.device("cuda"iftorch.cuda.is_available()else"cpu")dtype=torch.bfloat16batch_size,num_heads,head_dim=1,32,128max_seq_length=2048current_position=100# Create query and KV cacheq=torch.randn(batch_size,num_heads,1,head_dim,device=device,dtype=dtype)k_cache=torch.zeros(batch_size,num_heads,max_seq_length,head_dim,device=device,dtype=dtype)v_cache=torch.zeros(batch_size,num_heads,max_seq_length,head_dim,device=device,dtype=dtype)# Fill only valid positions (0-99)k_cache[:, :, :current_position]=torch.randn(batch_size,num_heads,current_position,head_dim,device=device,dtype=dtype)v_cache[:, :, :current_position]=torch.randn(batch_size,num_heads,current_position,head_dim,device=device,dtype=dtype)# Test 1: Current GPT-Fast approach (full cache)defoffset_causal_mask(b,h,q,kv):return (q+current_position-1)>=kvmask_full=create_block_mask(offset_causal_mask,B=batch_size,H=None,Q_LEN=1,KV_LEN=max_seq_length,device=device)mask_full.seq_lengths= (1,max_seq_length)# As done in generate.pyoutput_full=flex_attention(q,k_cache,v_cache,block_mask=mask_full)# Test 2: Correct approach (sliced cache)k_sliced=k_cache[:, :, :current_position]v_sliced=v_cache[:, :, :current_position]mask_sliced=create_block_mask(causal_mask,B=batch_size,H=None,Q_LEN=1,KV_LEN=current_position,device=device)mask_sliced.seq_lengths= (1,current_position)output_sliced=flex_attention(q,k_sliced,v_sliced,block_mask=mask_sliced)# Compare resultserror= (output_full-output_sliced).abs()print(f"Mean error:{error.mean().item():.6f}")print(f"Relative error:{(error.mean()/output_sliced.abs().mean()*100).item():.1f}%")print(f"Full cache std:{output_full.std().item():.6f}")print(f"Sliced cache std:{output_sliced.std().item():.6f}")

Mean error: 0.816406Relative error: 101.0%Full cache std: 0.152802Sliced cache std: 1.016770