@@ -769,10 +769,15 @@ def _gen_unused_masks(padding_mask, add_unused, max_seq_len, bs, device):
769769 # num_splits_vals = [1, 3]
770770 # SplitKV is not supported for hdim >= 192
771771 num_splits_vals = [1 , 3 ] if d < 192 and not DISABLE_SPLIT and not TEST_BWD_ONLY else [1 ]
772- for pack_gqa , num_splits in itertools .product (pack_gqa_vals , num_splits_vals ):
772+ precompute_metadata_vals = [False , True ]
773+ for pack_gqa , num_splits , precompute_metadata in itertools .product (
774+ pack_gqa_vals , num_splits_vals , precompute_metadata_vals
775+ ):
773776 # SplitKV not supported on SM90 - skip this iteration
774777 if IS_SM90 and num_splits > 1 :
775778 continue
779+ if precompute_metadata and is_fake_mode ():
780+ continue
776781 # TODO(wangsiyu): SM100 head_dim=256 2CTA kernel does not support pack_gqa yet.
777782 # pack_gqa=None means auto-enable for GQA/MQA (qhead_per_kvhead > 1)
778783 # Remove this when support is added.
@@ -781,56 +786,79 @@ def _gen_unused_masks(padding_mask, add_unused, max_seq_len, bs, device):
781786 continue
782787 if pack_gqa is None and mha_type != "mha" :
783788 continue
784- out_unpad , lse = flash_attn_varlen_func (
785- q_unpad if unpad_q else q ,
786- k_unpad if unpad_kv else k ,
787- v_unpad if unpad_kv else v ,
788- cu_seqlens_q = cu_seqlens_q if unpad_q else None ,
789- cu_seqlens_k = cu_seqlens_k if unpad_kv else None ,
790- max_seqlen_q = seqlen_q ,
791- max_seqlen_k = seqlen_k ,
792- seqused_q = seqused_q if not unpad_q else None ,
793- seqused_k = seqused_k if not unpad_kv else None ,
794- causal = causal ,
795- # qv=qv_unpad,
796- # q_descale=q_descale,
797- # k_descale=k_descale, v_descale=v_descale,
798- window_size = window_size ,
799- # attention_chunk=attention_chunk,
800- learnable_sink = learnable_sink ,
801- softcap = softcap ,
802- num_splits = num_splits ,
803- pack_gqa = pack_gqa ,
804- deterministic = deterministic ,
805- )
806- out = output_pad_fn (out_unpad ) if unpad_q else out_unpad
807- if is_fake_mode ():
808- # no more flash_attn cutedsl calls for the rest of the loop
809- # skip data-dependent postprocessing
810- continue
811- if query_unused_mask is not None :
812- out .masked_fill_ (q_zero_masking , 0.0 )
813- # When unpad_q=False with seqused_q, the kernel doesn't write positions
814- # beyond seqused_q, so those contain uninitialized values. Mask them out
815- # before comparing.
816- out_cmp , out_ref_cmp , out_pt_cmp = out , out_ref , out_pt
817- if not unpad_q and seqused_q is not None :
818- seqused_mask = torch .arange (seqlen_q , device = device )[None , :] < seqused_q [:, None ]
819- seqused_mask = rearrange (seqused_mask , "b s -> b s 1 1" )
820- out_cmp = out .clone ().masked_fill_ (~ seqused_mask , 0.0 )
821- out_ref_cmp = out_ref .clone ().masked_fill_ (~ seqused_mask , 0.0 )
822- out_pt_cmp = out_pt .clone ().masked_fill_ (~ seqused_mask , 0.0 )
823- print (f"Output max diff: { (out_cmp - out_ref_cmp ).abs ().max ().item ()} )
824- print (f"Output mean diff: { (out_cmp - out_ref_cmp ).abs ().mean ().item ()} )
825- # if not causal:
826- # print(f"LSE max diff: {(lse - lse_ref).abs().max().item()}")
827- # breakpoint()
789+ if precompute_metadata :
790+ scheduler_metadata = get_scheduler_metadata (
791+ num_batch = batch_size ,
792+ max_seqlen_q = seqlen_q ,
793+ max_seqlen_k = seqlen_k ,
794+ nheads = nheads ,
795+ nheads_kv = nheads_kv ,
796+ headdim = d ,
797+ headdim_v = dv ,
798+ num_splits = num_splits ,
799+ tile_m = 128 ,
800+ tile_n = 128 ,
801+ causal = causal ,
802+ cu_seqlens_q = cu_seqlens_q if unpad_q else None ,
803+ cu_seqlens_k = cu_seqlens_k if unpad_kv else None ,
804+ seqused_q = seqused_q if not unpad_q else None ,
805+ seqused_k = seqused_k if not unpad_kv else None ,
806+ )
807+ else :
808+ scheduler_metadata = None
809+ # Repeat to exercise metadata reuse across calls.
810+ for _ in range (1 if not precompute_metadata else 2 ):
811+ out_unpad , lse = flash_attn_varlen_func (
812+ q_unpad if unpad_q else q ,
813+ k_unpad if unpad_kv else k ,
814+ v_unpad if unpad_kv else v ,
815+ cu_seqlens_q = cu_seqlens_q if unpad_q else None ,
816+ cu_seqlens_k = cu_seqlens_k if unpad_kv else None ,
817+ max_seqlen_q = seqlen_q ,
818+ max_seqlen_k = seqlen_k ,
819+ seqused_q = seqused_q if not unpad_q else None ,
820+ seqused_k = seqused_k if not unpad_kv else None ,
821+ causal = causal ,
822+ # qv=qv_unpad,
823+ # q_descale=q_descale,
824+ # k_descale=k_descale, v_descale=v_descale,
825+ window_size = window_size ,
826+ # attention_chunk=attention_chunk,
827+ learnable_sink = learnable_sink ,
828+ softcap = softcap ,
829+ scheduler_metadata = scheduler_metadata ,
830+ num_splits = num_splits ,
831+ pack_gqa = pack_gqa ,
832+ deterministic = deterministic ,
833+ )
834+ out = output_pad_fn (out_unpad ) if unpad_q else out_unpad
835+ if is_fake_mode ():
836+ # no more flash_attn cutedsl calls for the rest of the loop
837+ # skip data-dependent postprocessing
838+ continue
839+ if query_unused_mask is not None :
840+ out .masked_fill_ (q_zero_masking , 0.0 )
841+ # When unpad_q=False with seqused_q, the kernel doesn't write positions
842+ # beyond seqused_q, so those contain uninitialized values. Mask them out
843+ # before comparing.
844+ out_cmp , out_ref_cmp , out_pt_cmp = out , out_ref , out_pt
845+ if not unpad_q and seqused_q is not None :
846+ seqused_mask = torch .arange (seqlen_q , device = device )[None , :] < seqused_q [:, None ]
847+ seqused_mask = rearrange (seqused_mask , "b s -> b s 1 1" )
848+ out_cmp = out .clone ().masked_fill_ (~ seqused_mask , 0.0 )
849+ out_ref_cmp = out_ref .clone ().masked_fill_ (~ seqused_mask , 0.0 )
850+ out_pt_cmp = out_pt .clone ().masked_fill_ (~ seqused_mask , 0.0 )
851+ print (f"Output max diff: { (out_cmp - out_ref_cmp ).abs ().max ().item ()} )
852+ print (f"Output mean diff: { (out_cmp - out_ref_cmp ).abs ().mean ().item ()} )
853+ # if not causal:
854+ # print(f"LSE max diff: {(lse - lse_ref).abs().max().item()}")
855+ # breakpoint()
828856
829- # Check that FlashAttention's numerical error is at most 3x the numerical error
830- # of a Pytorch implementation.
831- assert (out_cmp - out_ref_cmp ).abs ().max ().item () <= rtol * (
832- out_pt_cmp - out_ref_cmp
833- ).abs ().max ().item () + fwd_atol
857+ # Check that FlashAttention's numerical error is at most 3x the numerical error
858+ # of a Pytorch implementation.
859+ assert (out_cmp - out_ref_cmp ).abs ().max ().item () <= rtol * (
860+ out_pt_cmp - out_ref_cmp
861+ ).abs ().max ().item () + fwd_atol
834862
835863 if (
836864 dtype != torch .float8_e4m3fn
0 commit comments