fix unfused padding causal sdpa

[hungryGeek16]

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.

[greptile-apps]

Greptile Summary

This PR avoids materializing a full quadratic [b, 1, sq, sk] attention mask for padding_causal unfused attention by routing through a new per-sample F.scaled_dot_product_attention(is_causal=True) loop (_forward_varlen_sdpa) when the conditions are met. It also carries several unrelated improvements: MXFP8 attention backend filtering in get_attention_backend, FP8EmulationFunc null-quantizer guards, MXFP8 documentation updates, a new cp_size field in AttentionParams, and cuda_version added to the attention run-config log.

• Unfused padding_causal fast path: _use_varlen_sdpa gates a new _forward_varlen_sdpa method that iterates over batch items and calls PyTorch SDPA with is_causal=True, entirely bypassing get_full_mask. A new test (test_unfused_thd_padding_causal_uses_sdpa_without_full_mask) patches get_full_mask to assert it is never called and checks numerical correctness.
• MXFP8 FP8-emulation fixes: FP8EmulationFunc now handles None quantizers in the S_quantizer/O_quantizer/dO_quantizer/dP_quantizer branches and applies a BSHD→SBHD permute when the quantizer is an MXFP8Quantizer.

Confidence Score: 3/5

The new fast path produces correct results under normal configuration, but silently computes wrong attention scores when NVTE_APPLY_QK_LAYER_SCALING=1 is in use.

The new varlen-SDPA branch passes self.softmax_scale to _forward_varlen_sdpa instead of the locally-modified scale variable. Any deployment with NVTE_APPLY_QK_LAYER_SCALING=1 and padding_causal self-attention will silently receive incorrect attention output from the new fast path, with no error or warning.

transformer_engine/pytorch/attention/dot_product_attention/backends.py — specifically the scale argument at the _forward_varlen_sdpa call site.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/attention/dot_product_attention/backends.py	Adds _use_varlen_sdpa / _forward_varlen_sdpa fast path for padding_causal; passes self.softmax_scale instead of the locally-modified scale variable, silently breaking apply_qk_layer_scaling.
transformer_engine/pytorch/attention/dot_product_attention/utils.py	Adds MXFP8 backend filtering, cp_size field to AttentionParams, cuda_version to run_config, and refactors FA3/FA4 SM90 preference logic; changes appear correct.
tests/pytorch/attention/test_attention.py	Adds test_unfused_thd_padding_causal_uses_sdpa_without_full_mask which verifies both correctness and that get_full_mask is not called on the new fast path.
transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py	Documentation updates for MXFP8 recipe combinations; no functional logic changes.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[UnfusedDotProductAttention.forward] --> B{Convert input to sbhd}
    B --> C{padding in mask type and mask is None?}
    C -->|Yes| D[get_padding_mask: build mask from cu_seqlens]
    C -->|No| E[use existing mask]
    D --> F[_use_varlen_sdpa?]
    E --> F
    F -->|YES: padding_causal, self-attn, no bias, no dropout, no fp8, no alibi| G[_forward_varlen_sdpa: batch loop with SDPA is_causal=True, no full mask allocation]
    F -->|NO| H[get_full_mask: materialize full b x sq x sk mask]
    G --> I[_format_context: convert b,h,sq,d to output layout]
    H --> J[FP8 emulation if needed]
    J --> K[baddbmm QK matmul]
    K --> L[FusedScaleMaskSoftmax]
    L --> M[dropout + AV bmm]
    M --> N[reshape to output format]
    I --> Z[output tensor]
    N --> Z

Loading

_{Reviews (1): Last reviewed commit: "Avoid full mask allocation in unfused pa..." | Re-trigger Greptile}

[greptile-apps]

Wrong scale passed to fast-path SDPA

self.softmax_scale is passed here, but the locally-modified scale variable (which already incorporates the apply_qk_layer_scaling division by layer_number) was computed just a few lines above and is the value that should be forwarded. When NVTE_APPLY_QK_LAYER_SCALING=1 is set, scale = self.softmax_scale / self.layer_number, but the new varlen-SDPA path silently uses the unscaled value, producing wrong attention scores for every call that hits this branch with that env var enabled.

Suggested change

	return self._forward_varlen_sdpa(
	query_layer,
	key_layer,
	value_layer,
	q_format,
	batch_size,
	max_seqlen_q,
	cu_seqlens_q,
	attention_mask,
	self.softmax_scale,
	)
	if self._use_varlen_sdpa(
	attn_mask_type,
	attention_mask,
	window_size,
	core_attention_bias_type,
	alibi_slopes,
	fp8,
	):
	return self._forward_varlen_sdpa(
	query_layer,
	key_layer,
	value_layer,
	q_format,
	batch_size,
	max_seqlen_q,
	cu_seqlens_q,
	attention_mask,
	scale,
	)

[greptile-apps]

The first attn_mask_type != "padding_causal" guard already returns False for any value other than "padding_causal", so the later if attn_mask_type == "padding_causal" and … condition is always true at that point — the attn_mask_type == part is dead code and creates a misleading impression that the branch could be False.

Suggested change

	if attn_mask_type != "padding_causal":
	return False
	if window_size not in [None, (-1, 0), (-1, -1)]:
	return False
	if attn_mask_type == "padding_causal" and attention_mask is None:
	return False
	if attn_mask_type != "padding_causal":
	return False
	if window_size not in [None, (-1, 0), (-1, -1)]:
	return False
	if attention_mask is None:
	return False

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!