[QST][Cutlass Example 77] Softmax-corr Pipeline synchronization issue

[Aya-ZIbra]

Title: Pipeline synchronization issue in Flash Attention - S1 pipeline held during O0 rescale

Description:

While working on Flash Attention implementation using Cutlass pipelines and I'm encountering a synchronization issue that I don't understand.

In the correction function, when performing rescale(O0), the S1 pipeline appears to be held (not released) until after the O0 rescale completes. Same for S0 and O1 rescale. This seems to create unnecessary pipeline stalls, specially in causal/local masks.

Code context:

// After consuming S1 pipeline state 
pipeline_s1_c.consumer_wait(pipeline_s1_c_consumer_state);  

correction_rescale(scale, uint32_t(TmemAllocation::O1));
....
// During O0 rescale - S1 is NOT released here 
// pipeline_s1_c.consumer_release(pipeline_s1_c_consumer_state);  // Commented out 
// ++pipeline_s1_c_consumer_state;  

// S1 only gets released after O0 rescale 
correction_rescale(scale, uint32_t(TmemAllocation::O0)); 
pipeline_s1_c.consumer_release(pipeline_s1_c_consumer_state);  
++pipeline_s1_c_consumer_state;

Questions:

1. Why can't S1 be released immediately after it's consumed, before the O0 rescale?
2. Is there a dependency between S1 and O0 that I'm missing?

[hwu36]

@richardmcai

[hwu36]

@v0i0

[v0i0]

Hey Aya! Are you saying that you see speedups when you reorder those pipes?

To be honest I don't fully remember why it is written like that, so this is more speculation on my end, but here it goes:
In the softmax & correction groups, we respectively do the following (in the steady state):

def softmax():
    wait_for_qk_tc() # A
    compute_stats() # B
    write_stats() # C
    notify_correction_pipe() # D
    do_softmax_and_write_tmem() # E
    notify_mma_warp() # F
    reaquire_correction_pipe() # G
    update_local_sum() # H

def correction():
    acquire_s0() # A
   do_rescale() # B
   notify_s1() # C
   acquire_s1() # D
   do_rescale() # E
   notify_s0() # F

this implies the following pipeline:

softmax0    ABCDEF           GH ABCDEF           GH
softmax1           GH ABCDEF           GH ABCDEF
correction      ABC       DEF       ABC       DEF

which suggests to me that the purpose is to make sure step H / update_local_sum execute at a point where rescale is not executing.

[github-actions]

This issue has been labeled inactive-30d due to no recent activity in the past 30 days. Please close this issue if no further response or action is needed. Otherwise, please respond with a comment indicating any updates or changes to the original issue and/or confirm this issue still needs to be addressed. This issue will be labeled inactive-90d if there is no activity in the next 60 days.

[github-actions]

This issue has been labeled inactive-90d due to no recent activity in the past 90 days. Please close this issue if no further response or action is needed. Otherwise, please respond with a comment indicating any updates or changes to the original issue and/or confirm this issue still needs to be addressed.