support sparse gqa and fix flashattn

[STwangyingrui]

No description provided.

[gemini-code-assist]

Code Review

This pull request updates flash attention imports and reshaping logic, removes redundant device transfers for sequence lengths, and adds Grouped Query Attention (GQA) support in utility functions. Feedback was provided regarding the redundant recomputation and device transfer of cumulative sequence lengths in the transformer inference loop, suggesting they be moved outside the loop and explicitly cast to int32 for better performance and compatibility.

[gemini-code-assist]

The cu_seqlens_q and cu_seqlens_k tensors are recomputed and moved to the device for every layer in the transformer. Since query_lens and the updated key_values_lens are constant across all layers within a single inference step, these cumulative sequence lengths should ideally be computed once outside the layer loop (e.g., in the infer method) and passed down to self_attn to avoid redundant computations and expensive host-to-device transfers.

Additionally, you can combine the device move and type conversion into a single .to(device=AI_DEVICE, dtype=torch.int32) call to improve efficiency, as flash_attn requires int32 tensors.

Suggested change

	cu_seqlens_q = torch.nn.functional.pad(torch.cumsum(query_lens, dim=0), (1, 0)).to(AI_DEVICE)
	cu_seqlens_k = torch.nn.functional.pad(torch.cumsum(key_values_lens, dim=0), (1, 0)).to(AI_DEVICE)
	cu_seqlens_q = torch.nn.functional.pad(torch.cumsum(query_lens, dim=0), (1, 0)).to(device=AI_DEVICE, dtype=torch.int32)
	cu_seqlens_k = torch.nn.functional.pad(torch.cumsum(key_values_lens, dim=0), (1, 0)).to(device=AI_DEVICE, dtype=torch.int32)