Question about importing AdaSplash kernels in torchtitan - missing modules

[xiangguan1]

Dear developers,

First of all, thank you for your great work on AdaSplash! I'm very excited to try it out in my projects.

I'm trying to integrate AdaSplash with torchtitan following the instructions in your code. In attention.py, I see the following imports:

# AdaSplash kernels
from torchtitan.models.kernels.adaprefill_left_pad import sparse_attention_prefill
from torchtitan.models.kernels.adasplash15 import sparse_attn as adasplash15
from torchtitan.models.kernels.triton_entmax import triton_entmax_attention
from torchtitan.models.attention_bias import get_attn_mask_for_sdpa, get_slopes

However, I noticed that these modules don't seem to exist in the current torchtitan repository. When I try to run the code, I get a ModuleNotFoundError for adasplash.kernels (which appears to be related).

I was wondering:

1. Are these custom functions that I need to implement myself?
2. Or could this be a version compatibility issue between AdaSplash and torchtitan?
3. If these are part of AdaSplash, should I be importing them directly from the AdaSplash package instead of through torchtitan?

I've already installed AdaSplash via pip (adasplash==0.1.1b0), but I'm still encountering the missing module error when trying to use the attention integration.

Any guidance on how to properly set up and use AdaSplash with torchtitan would be greatly appreciated!

Thank you for your time and help!