Skip to content

[feat] QAD 5090: Wire the Attn-QAT inference attention backend (8/12) #1457

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
SolitaryThinker merged 2 commits into from
Jun 16, 2026
Merged

[feat] QAD 5090: Wire the Attn-QAT inference attention backend (8/12) #1457

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
15687b7
[feat]: Wire the Attn-QAT inference attention backend (8/12)
alexzms Jun 12, 2026
8996245
[bugfix]: plumb sm_scale through sageattn_blackwell + blockscaled_fp4…
SolitaryThinker Jun 16, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 9 additions & 5 deletions fastvideo-kernel/attn_qat_infer/api.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -135,13 +135,14 @@ def blockscaled_fp4_attn(qlist: Tuple,
is_causal: bool = False,
per_block_mean: bool = True,
is_bf16: bool = True,
single_level_p_quant: bool = False
single_level_p_quant: bool = False,
sm_scale: float | None = None
):
softmax_scale = (qlist[0].shape[-1] * 2) ** (-0.5)
softmax_scale = sm_scale if sm_scale is not None else (qlist[0].shape[-1] * 2) ** (-0.5)
return fp4attn_cuda.fwd(qlist[0], klist[0], vlist[0], qlist[1], klist[1], vlist[1], delta_s, KL, None, softmax_scale, is_causal, per_block_mean, is_bf16, single_level_p_quant)


def sageattn_blackwell(q, k, v, attn_mask = None, is_causal = False, per_block_mean = True, single_level_p_quant = True, **kwargs):
def sageattn_blackwell(q, k, v, attn_mask = None, is_causal = False, per_block_mean = True, single_level_p_quant = True, sm_scale: float | None = None, **kwargs):
"""
SageAttention3 Blackwell kernel for FP4 attention.

Expand All @@ -156,6 +157,8 @@ def sageattn_blackwell(q, k, v, attn_mask = None, is_causal = False, per_block_m
(standard per-block FP4 quantization like V, no s_P1).
If False (default), use two-level quantization:
s_P1 = rowmax(P̃)/(448×6), then s_P2, P̂_2 = φ(P̃/s_P1).
sm_scale: Softmax scale to pass through to the CUDA kernel. If None,
defaults to the kernel's 1/sqrt(D) scale.
**kwargs: Additional arguments (ignored)

Returns:
Expand All @@ -180,6 +183,7 @@ def sageattn_blackwell(q, k, v, attn_mask = None, is_causal = False, per_block_m
is_causal,
per_block_mean,
is_bf16,
single_level_p_quant
single_level_p_quant,
sm_scale
)[0][:, :, :QL, :].contiguous()
return o_fp4
return o_fp4
3 changes: 2 additions & 1 deletion fastvideo/configs/models/dits/base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,8 @@ class DiTArchConfig(ArchConfig):
AttentionBackendEnum.TORCH_SDPA,
AttentionBackendEnum.VIDEO_SPARSE_ATTN,
AttentionBackendEnum.VMOBA_ATTN, AttentionBackendEnum.SAGE_ATTN_THREE,
AttentionBackendEnum.SLA_ATTN, AttentionBackendEnum.SAGE_SLA_ATTN)
AttentionBackendEnum.ATTN_QAT_INFER, AttentionBackendEnum.SLA_ATTN,
AttentionBackendEnum.SAGE_SLA_ATTN)

hidden_size: int = 0
num_attention_heads: int = 0
Expand Down
7 changes: 7 additions & 0 deletions fastvideo/platforms/cuda.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -140,6 +140,13 @@ def get_attn_backend_cls(cls, selected_backend: AttentionBackendEnum | None, hea
except ImportError as e:
logger.info(e)
logger.info("Sage Attention 3 backend is not installed. Fall back to Flash Attention.")
elif selected_backend == AttentionBackendEnum.ATTN_QAT_INFER:
from fastvideo.attention.backends.attn_qat_infer import ( # noqa: F401
AttnQatInferBackend, is_attn_qat_infer_available)
if is_attn_qat_infer_available():
logger.info("Using Attn-QAT inference (modified SageAttention3 FP4) backend.")
return "fastvideo.attention.backends.attn_qat_infer.AttnQatInferBackend"
logger.info("Attn-QAT inference kernel is not built. Fall back to Flash Attention.")
Comment on lines +144 to +149

@gemini-code-assist gemini-code-assist Bot Jun 12, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

To improve robustness, wrap the import and availability check of the ATTN_QAT_INFER backend in a try...except block. This prevents potential runtime crashes (such as AttributeError if the kernel is missing expected attributes, or OSError due to CUDA driver/toolkit mismatches) and ensures a graceful fallback to Flash Attention. Additionally, we can remove the unused AttnQatInferBackend import and its associated # noqa: F401 suppression.

            try:
                from fastvideo.attention.backends.attn_qat_infer import is_attn_qat_infer_available
                if is_attn_qat_infer_available():
                    logger.info("Using Attn-QAT inference (modified SageAttention3 FP4) backend.")
                    return "fastvideo.attention.backends.attn_qat_infer.AttnQatInferBackend"
            except Exception as e:
                logger.info(e)
            logger.info("Attn-QAT inference kernel is not built. Fall back to Flash Attention.")

elif selected_backend == AttentionBackendEnum.VIDEO_SPARSE_ATTN:
try:
from fastvideo_kernel import video_sparse_attn # noqa: F401
Expand Down
1 change: 1 addition & 0 deletions fastvideo/platforms/interface.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ class AttentionBackendEnum(enum.Enum):
TORCH_SDPA = enum.auto()
SAGE_ATTN = enum.auto()
SAGE_ATTN_THREE = enum.auto()
ATTN_QAT_INFER = enum.auto()
VIDEO_SPARSE_ATTN = enum.auto()
BSA_ATTN = enum.auto()
VMOBA_ATTN = enum.auto()
Expand Down
Loading