rearrange(query, "b s h d -> (b s) h d") followed by batch_size, seq_len, num_head, hidden_dim = query.shape

[ulucsahin]

Hello and thank you for the great project.

I am trying to run wan_t2v_inference.sh, but I encountered the following issue.

in radial_attn/attn_mask.py line 327, query simple has 3 dimensions.

batch_size, seq_len, num_head, hidden_dim = query.shape

Before calling this, following rearrange is done:

radial-attention/radial_attn/models/wan/attention.py

Line 126 in d945317

query = rearrange(query, "b s h d -> (b s) h d")

and passed to radial attention, so it has 3 dimensions.

hidden_states = RadialAttention(
    query=query, key=key, value=value, mask_map=self.mask_map, sparsity_type="dense", block_size=128, decay_factor=self.decay_factor, model_type="wan", pre_defined_mask=None, use_sage_attention=self.use_sage_attention
)

If I simply set batch_size to 1 and do seq_len, num_head, hidden_dim = query.shape, it works but then I get errors afterwards.

[mali-afridi]

Do you happen to get this when trying to use FlashInferBackend instead of SpargeSageAttnBackend? I also did the same thing as you mentioned, but got this error afterwards:
RuntimeError: BatchPrefillWithPagedKVCacheSM90Run(at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor, std::optional<at::Tensor>, int64_t, int64_t, int64_t, double, double)::<lambda()>::<lambda()>::<lambda()>::<lambda()>::<lambda()> failed to dispatch (head_dim_qk, head_dim_vo): (40, 128)

[llw-nice]

I encountered the same issue. Have you resolved this problem?