Ordinary softmax divides each exponentiated score by the sum of exponentiated scores. For an all-masked row, every numerator is zero and the denominator is zero too. That is an undefined 0/0, often surfaced as NaN. Replacing -inf with a very large negative finite value may avoid the NaN, but it can assign probability to disallowed keys. That is not a safe fix. PyTorch's safe-softmax example demonstrates the all-masked behavior. Current PyTorch attention tutorials describe zero outputs for fully masked rows in supported implementations, so do not assume every backend still returns NaN. This question's custom path does.

I would capture the exact mask for the failing row, its sequence boundaries, padding lengths and whether the mask convention is True means allowed or True means blocked. Those conventions differ across APIs. Then compare the reference and fused kernel on the same tiny tensor, including forward output and backward gradient. If this row represents padding, it should usually produce a defined zero attention output and contribute no training loss. If it represents a real query token, an empty allowed-key set likely means the sequence metadata or causal mask is wrong. Masking the loss alone does not reliably contain NaNs that have already flowed through the network.

There is a tempting shortcut: always let the query attend to itself. That works only if the query is a real token and self-attention is permitted by the intended model. It can accidentally leak a padding token or break a cross-attention rule. Specify the semantics first. For packed data, test segment starts, zero-length segments if the pipeline permits them, left padding, right padding, and block boundaries. Ensure every real query has the intended allowed keys. For truly empty rows, use a safe kernel behavior and zero their output and gradient as required by the model contract.

Would I just drop the bad sample? For a one-off malformed record, perhaps, after identifying it. If the mask builder routinely creates empty rows, filtering a sample hides a structural bug that may also corrupt nonempty rows without NaNs. The tiled attention kernel returns NaNs. What did it forget when the row maximum changed? covers numerical rescaling when a tiled softmax maximum changes. Here the numerical implementation can be fine on valid rows. The input mask has no valid probability distribution for this query.