Handle left-padded batches without NaN logits

modeling_vortex.py

@@ -100,12 +100,21 @@ class VortexAttention(nn.Module):
             value = value.repeat_interleave(repeats, dim=1)
 
         sdpa_mask = None
+        query_valid = None
         is_causal = attention_mask is None and seq_len > 1
         if attention_mask is not None:
             if attention_mask.ndim == 2:
-                valid_keys = attention_mask.to(dtype=torch.bool, device=hidden_states.device)[:, None, None, :]
+                valid = attention_mask.to(dtype=torch.bool, device=hidden_states.device)
+                valid_keys = valid[:, None, None, :]
+                query_valid = valid[:, None, :, None]
                 causal = torch.ones((seq_len, seq_len), dtype=torch.bool, device=hidden_states.device).tril()
                 sdpa_mask = valid_keys & causal[None, None, :, :]
+                # SDPA returns NaN for an all-masked query row. Left padding
+                # creates exactly those rows, so give pad queries one harmless
+                # fallback key and zero their outputs after attention.
+                fallback = torch.zeros_like(sdpa_mask)
+                fallback[..., 0] = True
+                sdpa_mask = torch.where(query_valid, sdpa_mask, fallback)
             elif attention_mask.ndim == 4:
                 sdpa_mask = attention_mask
             else:
@@ -118,6 +127,8 @@ class VortexAttention(nn.Module):
             dropout_p=0.0,
             is_causal=is_causal,
         )
+        if query_valid is not None:
+            output = output * query_valid.to(dtype=output.dtype)
         output = output.transpose(1, 2).contiguous().view(batch, seq_len, -1)
         return self.o_proj(output)
 
@@ -230,7 +241,11 @@ class VortexForCausalLM(VortexPreTrainedModel):
         else:
             hidden_states = inputs_embeds
         if position_ids is None:
-            position_ids = torch.arange(hidden_states.shape[1], device=hidden_states.device).unsqueeze(0)
+            if attention_mask is not None and attention_mask.ndim == 2:
+                position_ids = attention_mask.long().cumsum(-1) - 1
+                position_ids = position_ids.masked_fill(attention_mask == 0, 0)
+            else:
+                position_ids = torch.arange(hidden_states.shape[1], device=hidden_states.device).unsqueeze(0)
         all_hidden_states = () if output_hidden_states else None
         for block in self.layers:
             if output_hidden_states: