def sparse_gqa_decode_varlen_indice(query, key, value, block_indices, cache_seqlens, max_cache_seqlen, block_size):
    """
    Args:
        query: [batch, heads, dim]
        key: [batch, max_cache_seqlen, heads_kv, dim]
        value: [batch, max_cache_seqlen, heads_kv, dim_v]
        block_indices: [batch, heads_kv, max_selected_blocks], indices of selected blocks, -1 for padding
        cache_seqlens: [batch], sequence lengths of the kvcache
        max_cache_seqlen: maximum sequence length of kvcache
        block_size: block size
    Returns:
        output: [batch, heads, dim_v]

    """
    
    batch, heads, dim = query.shape
    heads_kv = key.shape[2]
    dim_v = value.shape[-1]
    max_selected_blocks = block_indices.shape[-1]
    block_H = 64


    actual_num_blocks = torch.sum(block_indices != -1, dim=-1).to(torch.int32)
    actual_num_blocks = actual_num_blocks[:,0] #[batch],  number of valid blocks, assum all groups in the same batch have the same number of blocks
   
    # get num_split
    num_m_blocks = 1 * (heads // heads_kv + block_H - 1) // block_H
    num_n_blocks = max_selected_blocks#(kv_seqlen  + block_size - 1 ) // block_size
    # num_n_blocks = torch.sum(actual_num_blocks, dim=-1).item() * heads_kv # total number of blocks

    size_one_kv_head = max_selected_blocks * block_size * (dim + dim_v) * 2 #kv_seqlen * (dim + dim_v) * 2
    total_mblocks = batch * heads_kv * num_m_blocks
    num_sm = 132
    num_split = num_splits_heuristic(total_mblocks, num_sm, num_n_blocks, num_m_blocks, size_one_kv_head, is_causal_or_local=True, max_splits=128)
   
    program = flashattn(
        batch, heads, heads_kv, dim, dim_v)(
                block_N=block_size, block_H=block_H, num_split=T.symbolic("num_split"), num_stages=2, threads=128, 
                max_cache_seqlen=T.symbolic("max_cache_seqlen"), max_selected_blocks=T.symbolic("max_selected_blocks"))

    glse = torch.empty((batch, heads, num_split), dtype=torch.float32, device='cuda')
    Output_partial = torch.empty((batch, heads, num_split, dim_v), dtype=torch.float32, device='cuda')
    kernel = tilelang.compile(program, out_idx=-1, target='cuda', execution_backend="cython", pass_configs={"tl.config_index_bitwidth": 64})
    # print(kernel.get_kernel_source())

    # output = kernel(query, key, value, block_indices, cache_seqlens, actual_num_blocks, glse, Output_partial)
    output = kernel(query, key, value, block_indices, cache_seqlens, glse, Output_partial)
    return output

    Args:
        query: [batch, heads, dim]
        key: [batch, max_cache_seqlen, heads_kv, dim]
        value: [batch, max_cache_seqlen, heads_kv, dim_v]
        block_indices: [batch, heads_kv, max_selected_blocks], indices of selected blocks, -1 for padding
        cache_seqlens: [batch], sequence lengths of the kvcache
        max_cache_seqlen: maximum sequence length of kvcache
        block_size: block size
    Returns:
        output: [batch, heads, dim_v]

    
ReSA/llm/kernel/tilelang_sparse_decoding.py