Apply NSA to Qwen3

[Espere-1119-Song]

When I apply nsa to Qwen3, I meet the error AssertionError: Input shapes should have M >= 16, N >= 16 and K >= 16

I also use a demo to find the correct hypeparameter, but it always raises the same error. How can I fix it?

import torch
import sys
import os


sys.path.append(os.path.join(os.path.dirname(__file__), 'nsa-impl'))

from nsa.nsa import nsa_func

def test_nsa_with_small_dimensions():
    B, M, H, D = 1, 320, 64, 128  # Standard dimensions
    G = 16  # num_kv_heads = 16
    N = 320  # sequence length
    
    q = torch.randn(B, M, H, D, device='cuda')
    k = torch.randn(B, N, G, D, device='cuda')
    v = torch.randn(B, N, G, D, device='cuda')
    g_cmp = torch.randn(B, M, H, device='cuda')
    g_slc = torch.randn(B, M, H, device='cuda')
    g_swa = torch.randn(B, M, H, device='cuda')
    
    output = nsa_func(
        q, k, v,
        g_cmp=g_cmp,
        g_slc=g_slc,
        g_swa=g_swa,
        block_count=16,
        block_size=64,
        window_size=512,
        scale=None
    )
    


if __name__ == "__main__":
    test_nsa_with_small_dimensions() 

[mutiann]

I think it is due to the limitation of NSA that num_heads_q // num_heads_kv must >= 16. See https://www.zhihu.com/question/12616022631/answer/105896836750