Skip to content

custom_kernel: fix shape mismatch by sharding segment_ids in flash attn. #10397

custom_kernel: fix shape mismatch by sharding segment_ids in flash attn.

custom_kernel: fix shape mismatch by sharding segment_ids in flash attn. #10397