[Cute, Bwd, Sm100] Add varlen for sm100 bwd

[jayhshah]

We add varlen to sm100 backward pass and expose this capability through the flash_attn_varlen_func API.

Head-to-head benchmark against non-varlen sm100 backward for equal sequence lengths shows minimal overhead:

### headdim = 128, causal = False, seqlen = 8192, batch_size = 4, nheads = 16, nheads_kv = 16 ###
FA Python fwd: 1.445ms, 1522.0 TFLOPS
FA Python varlen fwd: 1.485ms, 1481.1 TFLOPS
FA Python bwd: 4.461ms, 1232.2 TFLOPS
FA Python varlen bwd: 4.577ms, 1201.0 TFLOPS
FA Python bwd (deterministic): 6.007ms, 915.1 TFLOPS
FA Python varlen bwd (deterministic): 6.020ms, 913.3 TFLOPS

### headdim = 128, causal = True, seqlen = 8192, batch_size = 4, nheads = 16, nheads_kv = 16 ###
FA Python fwd: 0.775ms, 1419.4 TFLOPS
FA Python varlen fwd: 0.792ms, 1388.1 TFLOPS
FA Python bwd: 2.447ms, 1123.4 TFLOPS
FA Python varlen bwd: 2.542ms, 1081.4 TFLOPS
FA Python bwd (deterministic): 2.923ms, 940.4 TFLOPS
FA Python varlen bwd (deterministic): 3.012ms, 912.6 TFLOPS

To fix an alignment issue with loading padded LSE in the backward kernel, we change padded offsets to FA3 style, e.g.:

padded_offset_q = seqlen.offset_q + batch_idx * self.m_block_size
if cutlass.const_expr(self.arch >= 90):
    padded_offset_q = padded_offset_q // self.m_block_size * self.m_block_size

[kiddyboots216]

varlen_fwd (and bwd) training matches FA2 on Blackwell

[v0i0]

is that meant to be an or? what is the logic here?

[jayhshah]

You're right to call this out, I only need to not use tma store dK/dV for cu_seqlens_k. I'll change this.

In general:

1. varlen k is the condition for the varlen scheduler
2. varlen q is a condition for checking process tile (since mblocks processed may then equal 0 if length for that batch is 0)

[jayhshah]

Though we also disable tma store for seqused_q only in the forward kernel, so there will be an inconsistency here (albeit with rarely used settings)

[jayhshah]

To address your other implicit question: note that since we use a special padded intermediate tensor for doing TMA reduce add with dK/dV accum when gqa, we are free to use tma store then without worrying about the usual problem of overwriting other batch's outputs. So it should be 'and' and not 'or'.

For the same reason we can use TMA reduce add for dQ accum.

[v0i0]

so the idea is: if we have gqa we post-process, so we can always use tma, even with seqlen_k. and we can use it without seqlen_k. so the only case where we need to not use tma is not seqlen_k and not gqa.

[jayhshah]

Yes (assuming you mean not tma = cu_seqlens_k and mha). I also tried using postprocess with mha and cu_seqlens_k to allow for tma store (hence why I had separated out the dKV_postprocess boolean) but that was slightly slower.