cleanup and add comments

Author: minosfuture

hopper/seqlen.h

@@ -56,7 +56,7 @@ struct SeqlenInfoQK {
                    ? seqlen_k_static
                    : (seqused_k ? seqused_k[bidb] : (cu_seqlens_k ? cu_seqlens_k[bidb + 1] - cu_seqlens_k[bidb] : seqlen_k_static)))
         , cp_world_size(cp_world_size)
-        , tot_seqlen_k(cp_tot_seqused_k == nullptr
+        , tot_seqlen_k(cp_tot_seqused_k == nullptr and cp_world_size <= 1
                        ? seqlen_k
                        : cp_tot_seqused_k[bidb])
     {

hopper/test_flash_attn.py

@@ -389,7 +389,7 @@ def test_flash_attn_output(
     ],
 )
 def test_flash_attn_varlen_output(
-        seqlen_q, seqlen_k, d, add_unused_qkv, causal, local, softcap, deterministic, has_qv_, mha_type, dtype, test_sink,
+        seqlen_q, seqlen_k, d, add_unused_qkv, causal, local, softcap, deterministic, has_qv_, mha_type, dtype, test_sink
 ):
     if has_qv_ and (d != 64 or dtype == torch.float8_e4m3fn):
         pytest.skip("Has Qv requires hdim 64 and dtype to be float16 or bfloat16 (not float8_e4m3fn)")

hopper/test_util.py

@@ -231,6 +231,7 @@ def construct_cp_mask(
         seqlen_k: Length of key sequence (local to this rank)
         cp_world_size: Number of context parallel ranks
         cp_rank: Current rank ID (0 to cp_world_size-1)
+        cp_tot_seqlen_k: Total lengths of key sequence in cp world
         window_size: (left_window, right_window), -1 = infinite
         sink_token_length: Number of "sink" tokens that can always be attended to
         query_padding_mask: Which query positions are valid
@@ -350,6 +351,7 @@ def attention_ref(
         s_aux: (nheads)
         cp_world_size: Number of context parallel ranks
         cp_rank: Current rank ID (0 to cp_world_size-1)
+        cp_tot_seqlen_k:  (batch_size) total seqlen of k/v in cp world
     Output:
         output: (batch_size, seqlen_q, nheads, head_dim_v)
         attention: (batch_size, nheads, seqlen_q, seqlen_k), softmax after dropout