Padding, Masks, and Variable-Length Batching

Padding produces a \(16 \times 1024\) tensor, so the attention score term is \(16 \times 1024^2 \approx 16.8\)M entries. The real requirement is \(1024^2 + 15 \times 32^2 \approx 1.06\)M, so roughly 94% of the work is on pad tokens, a factor of about 16. This is why varlen kernels exist: they iterate each sequence's own extent instead of a common rectangle.

A cumulative-length prefix array of size \(B+1\) marking each sequence's span in a flat token buffer; lengths \([1024, 32, 32]\) give \([0, 1024, 1056, 1088]\). Sequence \(i\) occupies the half-open range from entry \(i\) to entry \(i+1\). The two classic errors are omitting the trailing total, which truncates the last sequence, and passing a max_seqlen below the true maximum, which silently truncates the longest sequence's attention span.

Generation reads logits at the final index of the sequence tensor. With right padding, that index holds a pad token, so the model is asked to continue from <pad> rather than from the prompt. The output is fluent and wrong, and degrades gradually with padding amount rather than failing loudly. Position IDs must also be derived from the mask rather than arange, or every left-padded sequence starts at a nonzero offset.

Resetting position IDs at document boundaries. The mask correctly prevents cross-document attention, but RoPE still sees the second document beginning at index 1,024, teaching the model that a document can start anywhere in the position range. Masking and position indexing are independent pieces of plumbing and packing requires fixing both; getting only one right produces a subtle, quiet degradation.

Size. A materialised \(B \times H \times L \times L\) additive mask at \(B=8\), \(H=32\), \(L=8192\) in fp16 is roughly 34 GB, larger than the activations it guards, and reading it costs bandwidth on every attention call. A causal flag plus cu_seqlens describes the same constraint in a few kilobytes, and the kernel simply skips the blocks it would have masked rather than computing and then zeroing them.

Averaging loss within each microbatch and then averaging across microbatches gives every microbatch equal weight regardless of how many real tokens it held, so short sequences are effectively upweighted and the gradient depends on bucketing. Sum the per-token losses and divide once by the global count of unmasked tokens. The bug changes results without changing anything visible in the loss curve, since both versions descend smoothly.

Padding produces a \(16 \times 1024\) tensor, so the attention score term is \(16 \times 1024^2 \approx 16.8\)M entries. The real requirement is \(1024^2 + 15 \times 32^2 \approx 1.06\)M, so roughly 94% of the work is on pad tokens, a factor of about 16. This is why varlen kernels exist: they iterate each sequence's own extent instead of a common rectangle.

A cumulative-length prefix array of size \(B+1\) marking each sequence's span in a flat token buffer; lengths \([1024, 32, 32]\) give \([0, 1024, 1056, 1088]\). Sequence \(i\) occupies the half-open range from entry \(i\) to entry \(i+1\). The two classic errors are omitting the trailing total, which truncates the last sequence, and passing a max_seqlen below the true maximum, which silently truncates the longest sequence's attention span.

Generation reads logits at the final index of the sequence tensor. With right padding, that index holds a pad token, so the model is asked to continue from <pad> rather than from the prompt. The output is fluent and wrong, and degrades gradually with padding amount rather than failing loudly. Position IDs must also be derived from the mask rather than arange, or every left-padded sequence starts at a nonzero offset.

Resetting position IDs at document boundaries. The mask correctly prevents cross-document attention, but RoPE still sees the second document beginning at index 1,024, teaching the model that a document can start anywhere in the position range. Masking and position indexing are independent pieces of plumbing and packing requires fixing both; getting only one right produces a subtle, quiet degradation.

Size. A materialised \(B \times H \times L \times L\) additive mask at \(B=8\), \(H=32\), \(L=8192\) in fp16 is roughly 34 GB, larger than the activations it guards, and reading it costs bandwidth on every attention call. A causal flag plus cu_seqlens describes the same constraint in a few kilobytes, and the kernel simply skips the blocks it would have masked rather than computing and then zeroing them.

Averaging loss within each microbatch and then averaging across microbatches gives every microbatch equal weight regardless of how many real tokens it held, so short sequences are effectively upweighted and the gradient depends on bucketing. Sum the per-token losses and divide once by the global count of unmasked tokens. The bug changes results without changing anything visible in the loss curve, since both versions descend smoothly.