Uh oh!
There was an error while loading. Please reload this page.
fix unfused padding causal sdpa - #3063
Conversation
Greptile SummaryAdds a PyTorch SDPA fast path inside
Confidence Score: 3/5The fast path silently produces no gradients during training because context_layer is populated via in-place writes to a torch.zeros leaf tensor; this breaks autograd for every call that hits the new branch and needs to be resolved before merging. The _forward_varlen_sdpa method pre-allocates a zeros tensor and populates it with per-batch SDPA results through in-place indexing. PyTorch does not attach a grad_fn to a no-grad leaf tensor when it is the target of an in-place assignment, so the accumulated context_layer loses its connection to the computation graph. Every downstream op (permute, ConvertBSHDtoTHD, view) inherits the no-grad status, and the returned output has no grad_fn. In training this means query/key/value receive zero gradient updates on every call that takes the fast path. transformer_engine/pytorch/attention/dot_product_attention/backends.py — specifically _forward_varlen_sdpa and the context_layer accumulation pattern Important Files Changed
Flowchart%%{init: {'theme': 'neutral'}}%%
flowchart TD
A[forward called\nqkv_layout=thd_thd_thd] --> B[Convert THD → BSHD\nTranspose → sbhd]
B --> C{padding in attn_mask_type\nand mask is None?}
C -- yes --> D[get_padding_mask\nbatch,1,1,max_seqlen]
C -- no --> E[use provided mask]
D --> F[compute scale\napply_qk_layer_scaling]
E --> F
F --> G{_use_varlen_sdpa?\nattention_type==self\nattn_mask_type==padding_causal\nmax_sq==max_sk\nno alibi / fp8 / bias}
G -- True --> H[_forward_varlen_sdpa\nper-batch causal SDPA loop]
H --> I[context_layer = zeros⚠️\nno grad_fn]
I --> J[F.scaled_dot_product_attention\nis_causal=True per batch]
J --> K[in-place assign to zeros buffer\n⚠️ breaks autograd]
K --> L[_format_context\nConvertBSHDtoTHD for thd]
L --> M[return output\nrequires_grad=False ⚠️]
G -- False --> N[get_full_mask\nfull quadratic mask]
N --> O[regular unfused path\nbaddbmm + softmax + bmm]
O --> P[return output\nrequires_grad=True ✓]
Reviews (4): Last reviewed commit: "Avoid full mask allocation in unfused pa..." | Re-trigger Greptile |
Uh oh!
There was an error while loading. Please reload this page.
Uh oh!
There was an error while loading. Please reload this page.
cyanguwa
commented
Jun 1, 2026
Thanks for the contribution @hungryGeek16, but it looks like your base branch may be out of date - could you rebase please? Thanks! |
51d298b to
82f0e0eComparehungryGeek16
commented
Jun 8, 2026
@cyanguwa , I have rebased and resolved conflicts, let me know if this works. Thanks! |
| def _forward_varlen_sdpa( | ||
| self, | ||
| query_layer: torch.Tensor, | ||
| key_layer: torch.Tensor, | ||
| value_layer: torch.Tensor, | ||
| q_format: str, | ||
| batch_size: int, | ||
| max_seqlen_q: int, | ||
| cu_seqlens_q: Optional[torch.Tensor], | ||
| attention_mask: Optional[torch.Tensor], | ||
| scale: float, | ||
| ) -> torch.Tensor: | ||
| """Run causal self-attention without expanding padding masks to [b, 1, sq, sk].""" | ||
| context_layer = torch.zeros( | ||
| batch_size, | ||
| query_layer.size(2), | ||
| max_seqlen_q, | ||
| value_layer.size(3), | ||
| dtype=query_layer.dtype, | ||
| device=query_layer.device, | ||
| ) | ||
| if attention_mask is not None: | ||
| seqlens_q = attention_mask.logical_not()[:, 0, 0, :].sum(dim=1) | ||
| else: | ||
| seqlens_q = torch.full( | ||
| (batch_size,), max_seqlen_q, dtype=torch.int64, device=query_layer.device | ||
| ) | ||
| dropout_p = self.attention_dropout.p if self.training else 0.0 | ||
| with self.attention_dropout_ctx(): | ||
| for batch_id in range(batch_size): | ||
| seqlen_q = int(seqlens_q[batch_id].item()) | ||
| if seqlen_q == 0: | ||
| continue | ||
| query = query_layer[:seqlen_q, batch_id].permute(1, 0, 2).unsqueeze(0) | ||
| key = key_layer[:seqlen_q, batch_id].permute(1, 0, 2).unsqueeze(0) | ||
| value = value_layer[:seqlen_q, batch_id].permute(1, 0, 2).unsqueeze(0) | ||
| context_layer[batch_id, :, :seqlen_q, :] = F.scaled_dot_product_attention( | ||
| query, | ||
| key, | ||
| value, | ||
| dropout_p=dropout_p, | ||
| is_causal=True, | ||
| scale=scale, | ||
| ).squeeze(0) |
There was a problem hiding this comment.
_use_varlen_sdpa has no guard on max_seqlen_q == max_seqlen_kv. For inference with a KV cache (e.g. qkv_format = "sbhd_2bshd" or "thd_2bshd"), max_seqlen_q is the current decode length (often 1) while max_seqlen_kv is the full cache length. When padding_causal is set for batched inference, get_padding_mask creates a [batch, 1, 1, max_seqlen_q] mask, seqlens_q will be all-1s, and key_layer[:1, batch_id] selects only the first cache token instead of the full KV context. Every generated token then attends to nothing but the first token — silently wrong, no error raised.
The fix is to add max_seqlen_q == max_seqlen_kv as a guard in _use_varlen_sdpa (forwarding those values from forward), or to reject the fast path inside forward before calling _use_varlen_sdpa when the two dimensions differ. FlashAttention.forward already asserts this invariant (line 1054) for the same reason.
Signed-off-by: hungryGeek16 <rahul_mangalampalli@yahoo.in>
166a496 to
c27595fCompare| cu_seqlens_q: Optional[torch.Tensor], | ||
| attention_mask: Optional[torch.Tensor], | ||
| scale: float, | ||
| ) -> torch.Tensor: | ||
| """Run causal self-attention without expanding padding masks to [b, 1, sq, sk].""" | ||
| context_layer = torch.zeros( | ||
| batch_size, | ||
| query_layer.size(2), | ||
| max_seqlen_q, | ||
| value_layer.size(3), | ||
| dtype=query_layer.dtype, | ||
| device=query_layer.device, | ||
| ) | ||
| if attention_mask is not None: | ||
| seqlens_q = attention_mask.logical_not()[:, 0, 0, :].sum(dim=1) | ||
| else: | ||
| seqlens_q = torch.full( | ||
| (batch_size,), max_seqlen_q, dtype=torch.int64, device=query_layer.device | ||
| ) | ||
| dropout_p = self.attention_dropout.p if self.training else 0.0 | ||
| with self.attention_dropout_ctx(): | ||
| for batch_id in range(batch_size): | ||
| seqlen_q = int(seqlens_q[batch_id].item()) | ||
| if seqlen_q == 0: | ||
| continue | ||
| query = query_layer[:seqlen_q, batch_id].permute(1, 0, 2).unsqueeze(0) | ||
| key = key_layer[:seqlen_q, batch_id].permute(1, 0, 2).unsqueeze(0) | ||
| value = value_layer[:seqlen_q, batch_id].permute(1, 0, 2).unsqueeze(0) | ||
| context_layer[batch_id, :, :seqlen_q, :] = F.scaled_dot_product_attention( | ||
| query, | ||
| key, | ||
| value, | ||
| dropout_p=dropout_p, | ||
| is_causal=True, | ||
| scale=scale, | ||
| ).squeeze(0) |
There was a problem hiding this comment.
context_layer = torch.zeros(...) creates a leaf tensor with requires_grad=False. Each context_layer[batch_id, :, :seqlen_q, :] = F.scaled_dot_product_attention(...).squeeze(0) is an in-place write to that no-grad leaf. PyTorch does not attach a grad_fn to a no-grad leaf when it is the target of an in-place op, so context_layer retains requires_grad=False after all iterations. Every downstream operation in _format_context (permute, ConvertBSHDtoTHD, view) also produces no-grad tensors, and the final output from _forward_varlen_sdpa has no grad_fn.
At training time this means backward silently skips the entire fast path — gradients for query_layer, key_layer, and value_layer are never computed, so the attention projection weights receive zero updates. In the regression test, output.float().sum().backward() raises RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn (or leaves all .grad as None), causing the assertions on lines 726-729 to fail.
The fix is to avoid the pre-allocated zeros buffer and instead build context_layer from the SDPA outputs using differentiable stacking, e.g. collect per-batch padded results in a list and torch.cat them.
| dropout_p = self.attention_dropout.p if self.training else 0.0 | ||
| with self.attention_dropout_ctx(): | ||
| for batch_id in range(batch_size): |
There was a problem hiding this comment.
This seems to be serializing the sequences in a batch. Is this really more efficient than the original implementation in UnfusedDotProductAttention?
Also, is there any particular reason to use Torch SDPA in TE's unfused path, when THD padding_causal self-attention is supported by FusedAttention and FlashAttention? Thanks.
Adds a targeted PyTorch SDPA fallback for unfused THD padding_causal self-attention so TransformerEngine does not materialize the full quadratic padding/causal mask. Includes a regression test that fails if get_full_mask is called on this path.