flash_attn ¶

def __init__(
    self,
    num_heads: int,
    head_size: int,
    causal: bool,
    softmax_scale: float,
    num_kv_heads: int | None = None,
    prefix: str = "",
    **extra_impl_args,
) -> None:
    self.causal = causal
    self.softmax_scale = softmax_scale
    # MiniMax-H3's dense DiT explicitly enables this faster FA4 entry
    # point. It remains off for every other model and for the H3 text
    # refiner; grad-enabled calls stay on the established fixed path.
    self.fa4_packed_varlen = bool(extra_impl_args.get("fa4_packed_varlen", False))
    if self.fa4_packed_varlen and fa_version == "4":
        _log_fa4_packed_varlen_config()
    # An explicit ``nvfp4_fa4`` impl arg wins over the process-wide
    # FASTVIDEO_NVFP4_FA4 env opt-in, so precision-sensitive layers (e.g.
    # the FP32-pinned H3 VAE attention) can force-disable FP4 Q/K
    # quantization while the DiT keeps it. When the arg is absent the env
    # keeps its previous semantics.
    nvfp4_fa4 = extra_impl_args.get("nvfp4_fa4")
    if nvfp4_fa4 is None:
        nvfp4_fa4 = os.environ.get("FASTVIDEO_NVFP4_FA4", "0") == "1"
    self.nvfp4_fa4 = bool(nvfp4_fa4)
    if self.nvfp4_fa4:
        cap = torch.cuda.get_device_capability()
        assert cap in [(10, 0), (10, 3)], (f"NVFP4 FA4 requires Blackwell (sm100a/sm103a), got sm{cap[0]}{cap[1]}")
        assert _FA4_FP4_AVAILABLE, ("NVFP4 FA4 requires flash-attention-fp4 (flash_attn.cute). "
                                    "Install via instructions in docs/inference/optimizations.md")
        logger.info("NVFP4 FA4 enabled for FlashAttentionImpl (quant_qk only)")