[Cute,Fwd,Sm100] Add r2p for local mask

[henrylhtsang]

The idea is similar to mask_r2p.

Using xor mask_range = mask_right & ~ mask_left because:

upper bound: 0b00011111
lower bound: 0b00000001
upper & ~ lower : 0b00011110

Notes:

• Original idea was to use xor. But that would require assuming window_size_left + window_size_right >= 0. I can't tell if xor is really better.
• I wanted to avoid taking min(right_s, 24). But I would encounter an error OverflowError: Python int too large to convert to C long at mask_range = mask_right & ~mask_left

Benchmark

Before

### headdim = 128, causal = False, seqlen = 8192, batch_size = 4, nheads = 16, nheads_kv = 16, window=symmetric(512,512), varlen = False, deterministic = False ###
FA Python fwd: 0.304ms, 875.9 TFLOPS

### headdim = 128, causal = False, seqlen = 8192, batch_size = 4, nheads = 16, nheads_kv = 16, window=symmetric(1024,1024), varlen = False, deterministic = False ###
FA Python fwd: 0.442ms, 1167.5 TFLOPS

### headdim = 128, causal = False, seqlen = 8192, batch_size = 4, nheads = 16, nheads_kv = 16, window=symmetric(2048,2048), varlen = False, deterministic = False ###
FA Python fwd: 0.723ms, 1330.8 TFLOPS

### headdim = 128, causal = False, seqlen = 8192, batch_size = 4, nheads = 16, nheads_kv = 16, window=symmetric(4096,4096), varlen = False, deterministic = False ###
FA Python fwd: 1.144ms, 1442.1 TFLOPS

### headdim = 128, causal = False, seqlen = 8192, batch_size = 4, nheads = 16, nheads_kv = 16, window=left(512,0), varlen = False, deterministic = False ###
FA Python fwd: 0.232ms, 574.0 TFLOPS

### headdim = 128, causal = False, seqlen = 8192, batch_size = 4, nheads = 16, nheads_kv = 16, window=left(1024,0), varlen = False, deterministic = False ###
FA Python fwd: 0.297ms, 869.3 TFLOPS

### headdim = 128, causal = False, seqlen = 8192, batch_size = 4, nheads = 16, nheads_kv = 16, window=left(2048,0), varlen = False, deterministic = False ###
FA Python fwd: 0.417ms, 1152.8 TFLOPS

### headdim = 128, causal = False, seqlen = 8192, batch_size = 4, nheads = 16, nheads_kv = 16, window=left(4096,0), varlen = False, deterministic = False ###
FA Python fwd: 0.635ms, 1299.5 TFLOPS

After

### headdim = 128, causal = False, seqlen = 8192, batch_size = 4, nheads = 16, nheads_kv = 16, window=symmetric(512,512), varlen = False, deterministic = False ###
FA Python fwd: 0.281ms, 948.1 TFLOPS

### headdim = 128, causal = False, seqlen = 8192, batch_size = 4, nheads = 16, nheads_kv = 16, window=symmetric(1024,1024), varlen = False, deterministic = False ###
FA Python fwd: 0.426ms, 1209.6 TFLOPS

### headdim = 128, causal = False, seqlen = 8192, batch_size = 4, nheads = 16, nheads_kv = 16, window=symmetric(2048,2048), varlen = False, deterministic = False ###
FA Python fwd: 0.708ms, 1358.8 TFLOPS

### headdim = 128, causal = False, seqlen = 8192, batch_size = 4, nheads = 16, nheads_kv = 16, window=symmetric(4096,4096), varlen = False, deterministic = False ###
FA Python fwd: 1.131ms, 1458.5 TFLOPS

### headdim = 128, causal = False, seqlen = 8192, batch_size = 4, nheads = 16, nheads_kv = 16, window=left(512,0), varlen = False, deterministic = False ###
FA Python fwd: 0.206ms, 648.4 TFLOPS

### headdim = 128, causal = False, seqlen = 8192, batch_size = 4, nheads = 16, nheads_kv = 16, window=left(1024,0), varlen = False, deterministic = False ###
FA Python fwd: 0.275ms, 938.9 TFLOPS

### headdim = 128, causal = False, seqlen = 8192, batch_size = 4, nheads = 16, nheads_kv = 16, window=left(2048,0), varlen = False, deterministic = False ###
FA Python fwd: 0.400ms, 1202.3 TFLOPS

### headdim = 128, causal = False, seqlen = 8192, batch_size = 4, nheads = 16, nheads_kv = 16, window=left(4096,0), varlen = False, deterministic = False ###
FA Python fwd: 0.621ms, 1329.1 TFLOPS

test

pytest . -x 

============================================================================================ 49957 passed, 34769 skipped, 1325295 warnings in 20300.02s (5:38:20) ============================================================================================

[henrylhtsang]

maybe @jayhshah @v0i0 ?

[henrylhtsang]

I compared xor vs mask_range = mask_right & ~ mask_left. They are pretty close, probably 0.5% difference, but xor seems faster.

updated to use out_bound. perf is better. nah

[tridao]

This is awesome!

[henrylhtsang]

@v0i0 noob question: I don't see a merge button. Can you help merge for me? tests just finished running