Using the same selection algorithm, but employing non-quantized Flashinfer for computation, why are the generated results all garbled? #106

def spas_sage2_attn_meansim_topk_cuda(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None, smooth_k=True, simthreshd1=-0.1, cdfthreshd=None, topk=0.5, pvthreshd=50, attention_sink=False, tensor_layout="HND", output_dtype=torch.float16, return_sparsity=False):
 assert tensor_layout in ['HND', 'NHD']
 if tensor_layout == 'NHD':
 q, k, v = map(lambda t: rearrange(t, '... L H D -> ... H L D'), (q, k, v))
 assert q.size(-2)>=128, "seq_len should be not less than 128."
 torch.cuda.set_device(v.device)
 dtype = q.dtype
 if dtype == torch.float32 or dtype == torch.float16:
 q, k, v = q.contiguous().to(torch.float16), k.contiguous().to(torch.float16), v.contiguous().to(torch.float16)
 else:
 q, k, v = q.contiguous().to(torch.bfloat16), k.contiguous().to(torch.bfloat16), v.contiguous().to(torch.float16)
 if smooth_k:
 km = k.mean(dim=-2, keepdim=True)
 # k = k - km
 headdim = q.size(-1)
 arch = get_cuda_arch_versions()[q.device.index]
 if arch == "sm90":
 lut, valid_block_num, q_int8, q_scale, k_int8, k_scale = get_block_map_meansim_fuse_quant(q, k, km, is_causal=is_causal, simthreshd1=simthreshd1, cdfthreshd=cdfthreshd, topk=topk, return_lut=True, attention_sink=attention_sink, BLKQ=64, BLKK=128)
 else:
 lut, valid_block_num, q_int8, q_scale, k_int8, k_scale = get_block_map_meansim_fuse_quant(q, k, km, is_causal=is_causal, simthreshd1=simthreshd1, cdfthreshd=cdfthreshd, topk=topk, return_lut=True, attention_sink=attention_sink, BLKQ=128, BLKK=64)
 if scale is None:
 scale = 1.0 / (headdim ** 0.5)
 assert headdim in [64, 128], "headdim should be in [64, 128]. For other headdim, you can use padding and specify the softmax scale."
 pvthreshd = hyperparameter_check(pvthreshd, q.size(-3), q.device)
 ## quant v
 b, h_kv, kv_len, head_dim = v.shape
 padded_len = (kv_len + 127) // 128 * 128
 v_transposed_permutted = torch.empty((b, h_kv, head_dim, padded_len), dtype=v.dtype, device=v.device)
 fused.transpose_pad_permute_cuda(v, v_transposed_permutted, 1)
 v_fp8 = torch.empty(v_transposed_permutted.shape, dtype=torch.float8_e4m3fn, device=v.device)
 v_scale = torch.empty((b, h_kv, head_dim), dtype=torch.float32, device=v.device)
 #fused.scale_fuse_quant_cuda(v_transposed_permutted, v_fp8, v_scale, kv_len, 448.0, 1)
 fused.scale_fuse_quant_cuda(v_transposed_permutted, v_fp8, v_scale, kv_len, 2.25, 1)
 _is_causal = 1 if is_causal else 0
 o = torch.empty_like(q)
 
 if arch == "sm90":
 qattn.qk_int8_sv_f8_accum_f32_block_sparse_attn_inst_buf_fuse_v_scale_with_pv_threshold_sm90(q_int8, k_int8, v_fp8, o, lut, valid_block_num, pvthreshd, q_scale, k_scale, v_scale, 1, False, 1, scale, 0)
 elif SAGE2PP_ENABLED:
 qk_int8_sv_f8_accum_f16_block_sparse_attn_inst_buf_fuse_v_scale_with_pv_threshold(q_int8, k_int8, v_fp8, o, lut, valid_block_num, pvthreshd, q_scale, k_scale, v_scale, 1, False, 1, scale, 0)
 else:
 qattn.qk_int8_sv_f8_accum_f32_block_sparse_attn_inst_buf_fuse_v_scale_with_pv_threshold(q_int8, k_int8, v_fp8, o, lut, valid_block_num, pvthreshd, q_scale, k_scale, v_scale, 1, False, 1, scale, 0)
 
 if tensor_layout == 'NHD':
 o = rearrange(o, '... H L D -> ... L H D')
 return o

@torch.compiler.disable
def spas_sage2_attn_meansim_topk_cuda(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None, smooth_k=True, simthreshd1=-0.1, cdfthreshd=None, topk=0.5, pvthreshd=50, attention_sink=False, tensor_layout="HND", output_dtype=torch.float16, return_sparsity=False):
 block_q = 128
 block_k = 64
 assert tensor_layout in ['HND', 'NHD']
 torch.cuda.set_device(v.device)
 dtype = q.dtype
 if dtype == torch.float32 or dtype == torch.float16:
 q, k, v = q.contiguous().to(torch.float16), k.contiguous().to(torch.float16), v.contiguous().to(torch.float16)
 else:
 q, k, v = q.contiguous().to(torch.bfloat16), k.contiguous().to(torch.bfloat16), v.contiguous().to(torch.float16)
 if smooth_k:
 km = k.mean(dim=-2, keepdim=True)
 # k = k - km
 headdim = q.size(-1)
 arch = get_cuda_arch_versions()[q.device.index]
 final_map, _ , _ , _ , _ = get_block_map_meansim_fuse_quant(q, k, km, is_causal=is_causal, simthreshd1=simthreshd1, cdfthreshd=cdfthreshd, topk=topk, return_lut=False, attention_sink=attention_sink, BLKQ=block_q, BLKK=block_k)
 b = q.size(0)
 num_heads = q.size(1)
 seq_len = q.size(2)
 dim_head = q.size(3)
 final_map = final_map.reshape(b * num_heads, final_map.shape[2], final_map.shape[3])
 block_row_sz = torch.full((b * num_heads, final_map.shape[1]), block_q, dtype=torch.int32, device=q.device)
 block_column_sz = torch.full((b * num_heads, final_map.shape[2]), block_k, dtype=torch.int32, device=q.device) 
 q_permuted = q.reshape(b * num_heads, seq_len, dim_head).contiguous()
 k_permuted = k.reshape(b * num_heads, seq_len, dim_head).contiguous()
 v_permuted = v.reshape(b * num_heads, seq_len, dim_head).contiguous()
 bsr_wrapper.plan(
 final_map,
 block_row_sz,
 block_column_sz,
 b * num_heads, 
 b * num_heads, 
 dim_head, 
 sm_scale = 1.0 / (dim_head ** 0.5),
 q_data_type = q.dtype,
 kv_data_type = k.dtype,
 )
 out = bsr_wrapper.run( 
 q_permuted,
 k_permuted,
 v_permuted
 )
 torch.cuda.synchronize()
 out = out.reshape(b, num_heads, seq_len, dim_head)
 return out