keep sm90 headsize 128 cubins (#5320)

Signed-off-by: Qidi Sang <[email protected]>

Author: qsang-nv

cpp/kernels/fmha_v2/README.md

@@ -3049,14 +3049,20 @@ def get_kernel_traits_code(specs_names):
     return code
 
 
+# For now, only hopper head_size 128 kernel uses cubins, and other kernels use cu files.
+# You should set the condition `use_cubin_header` to false if you have modified the source code of the FMHA kernels on Hopper (sm90) with head_size 128.
+# This ensures that the kernels will be recompiled using the updated source code rather than relying on precompiled cubins.
+def use_cubin_header(kspec):
+    return kspec.sm == 90 and kspec.head_size == 128
+
+
 def get_cubin_header(kernel_traits, specs_names):
     cubins = []
     cubin_lens = []
     cubins_dict = {}
     cubin_lens_dict = {}
     for kspec, fname, lname, kname in specs_names:
-        # only generate hopper cubin header
-        if generate_cu_trtllm and not 'sm90' in kname:
+        if generate_cu_trtllm and not use_cubin_header(kspec):
             continue
         name = fname.replace('.', '_')
         data = 'extern unsigned char cubin_{name}_cubin[];'.format(name=name)
@@ -3209,7 +3215,7 @@ def get_cubin_header(kernel_traits, specs_names):
             if generate_cu_trtllm:
 
                 def get_lname_from_kname(kname: str) -> str:
-                    if 'sm90' in kname:
+                    if use_cubin_header(kspec):
                         return 'nullptr'
                     lname = kname.replace('_kernel', '')
                     mask_types = [
@@ -3228,7 +3234,7 @@ def get_lname_from_kname(kname: str) -> str:
 {cubin_name}_len, \"{kname}\", {smem}, {threads}, {meta_unroll_step}, {attention_mask_type_value}, \
 {attention_input_layout_value}, {is_il}, {is_flash_atten}, {is_warp_specialization}, {is_fp32_accu}, \
 {is_alibi_supported}, {is_tiled}, {has_softcapping_scale}, {return_softmax_stats_flag}, {lname}}}\
-'''.format(**locals()) if 'sm90' in kname else '''\
+'''.format(**locals()) if use_cubin_header(kspec) else '''\
 {{ DATA_TYPE_{prec}, DATA_TYPE_{output_prec}, {seq_len}, {q_step}, {kv_step}, {head_size}, {head_size_v}, \
 {sage_block_sizes[0]}, {sage_block_sizes[1]}, {sage_block_sizes[2]}, kSM_{sm}, nullptr, \
 0, \"{kname}\", {smem}, {threads}, {meta_unroll_step}, {attention_mask_type_value}, \
@@ -3404,6 +3410,9 @@ def get_lname_from_kname(kname: str) -> str:
     return code
 
 
+# This is used to add some kernels running in cubins.
+# The source code of paged context fmha kernels are not in this repo, but we have cubins for them.
+# Other kernels are for passing CI cases.
 def modify_cubin_header(cubin_header):
     # for paged context fmha cases
     target = "#ifndef EXCLUDE_SM_90"