fix - fix bug when not relevant seq has nan data

[baowendin]

We found that in the batch MLA FA2 implementation, if there exists NaN data within the computed CKV cache page, the kernel implementation would inadvertently read these invalid values and corrupt other valid data, ultimately resulting in all outputs becoming NaN. Therefore, this PR introduces masking logic to exclude CKV entries outside the computation scope from being processed. It may cause other requests to get error response when previous request precision overflow
and write nan to kvcache.

How to reproduce

1. modify test_batch_mla_varlen_page_attention in test_deepseek_mla.py to use fa2 implement, and set an nan data in unrelated sequence

    # disable fa3 check
    # if not is_sm90a_supported(torch.device("cuda")):
    #     pytest.skip("FA3 is not supported on this device")
    if causal and qo_len > min(kv_len_0, kv_len_1, kv_len_2):
        pytest.skip("qo_len > kv_len not supported for causal attention")
    num_different_kv_len = 3
    .....
    kpe = torch.randn(
        batch_size * pages_nums_sum,
        page_size,
        head_dim_kpe,
        dtype=dtype,
        device="cuda",
    )
    # set a not related seq to nan
    ckv[1][10] = float("nan")

2. use test in put as below

test_deepseek_mla.py

if __name__ == "__main__":    
    test_batch_mla_varlen_page_attention(
        1, 65, 65, 65, 1, 128, True, 64, "fa2", torch.half
    )

3. nan output can be observed in output

2025-03-14 10:06:40,470 - INFO - flashinfer.jit: Finished loading JIT ops: batch_mla_attention_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_ckv_512_head_dim_kpe_64_profiler_False
tensor([[[    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         ...,
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan]],

        [[ 0.2800, -0.1050,  0.0584,  ..., -0.3953,  0.2433,  0.0251],
         [ 0.2006,  0.2817, -0.1516,  ...,  0.2883,  0.2686, -0.1577],
         [ 0.1216, -0.8970, -0.1066,  ...,  0.5020, -0.7349,  0.3784],
         ...,
         [ 0.0591,  0.5034, -0.3750,  ..., -0.2786, -0.2483,  0.0037],
         [ 0.2854,  0.1653, -0.1549,  ..., -0.5571, -0.3015,  0.0625],
         [ 0.0942, -0.3574,  0.1144,  ...,  0.1730, -0.2646,  0.2479]],

        [[ 0.0208, -0.3357,  0.2020,  ..., -0.1625, -0.0917,  0.2349],
         [ 0.0799,  0.6021, -0.0385,  ...,  0.5962,  0.3550,  0.2208],
         [ 0.3367, -0.5327, -0.2496,  ...,  0.0653, -0.2144,  0.1873],
         ...,
         [ 0.0577, -1.1670,  1.5000,  ...,  1.4023,  0.7612, -0.6890],
         [-0.2124, -0.2549,  0.3342,  ..., -0.5610, -0.1438, -0.1825],
         [ 0.0499, -0.3467, -0.1461,  ..., -0.0058,  0.0474,  0.3188]]],
       device='cuda:0', dtype=torch.float16)

[yzh119]

I think there might be a easier way:

1. passing kv_end instead of last_page_kv_idx to this function.
2. change the condition to (packed_block_iter_base + lane_idx / 8 + warp_idx_in_wg * 4) < kv_end

WDYT?

[yzh119]

I created anther bugfix (#945) for FA3 template, does that look good to you? If so, I suppose we can use similar form here.

[baowendin]

okay, let me modify the commit

[yzh119]

sorry I directly modify the commit and merge it, let me know if you think there any issue with it