Fixed deadlock in sgmv_shrink kernel caused by skewed segments

[tgaddair]

When there is a large imbalance (>= 65 elements in the batch) in the size of two or more segments in a batch, it can lead to deadlocks in the sgmv_shrink kernel.

The crux of the issue was that each grid block can execute a dynamic number of steps depending on the size of its segment (s_end - s_start). However, during each step the block will call grid.sync(). If one block executes more steps than another, it will call grid.sync() a different number of times, leading to a deadlock.

The solution presented here is to compute the max number of steps from the largest segment, and then call grid.sync() at the end of the kernel for the difference between the max steps and the current block's steps.

Because the length of the s vector is generally very small (< batch size), the loop here should not introduce noticeable latency. However, it may be worth exploring more optimized solutions to this problem in a follow-up.

Note that this issue only occurs when using cooperative groups.

Related:

• Fixed deadlock in sgmv_shrink kernel caused by imbalanced segments predibase/lorax#156
• Lorax Hanging in production predibase/lorax#149

[tgaddair]

cc @abcdabcd987

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Comparison is base (07a40b9) 43.27% compared to head (2a32749) 43.27%.
Report is 1 commits behind head on master.

Additional details and impacted files

@@           Coverage Diff           @@
##           master      #35   +/-   ##
=======================================
  Coverage   43.27%   43.27%           
=======================================
  Files          10       10           
  Lines         647      647           
=======================================
  Hits          280      280           
  Misses        367      367           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[abcdabcd987]

Thanks!

@yzh119 Can you take a look?

[yzh119]

Thank you @tgaddair for the patch, yes I do think you are correct.