test mla and nsa performance on modal

[120L021326]

No description provided.

[120L021326]

mla_decoding log:

-------------------------------
Running on TestParam(b=132, s_q=2, s_k=4096, is_varlen=False, is_causal=False, is_fp8=False, topk=None, test_performance=True, is_all_indices_invalid=False, have_zero_seqlen_k=False, block_size=64, h_q=128, h_kv=1, d=576, dv=512, seed=0)...
Correctness check passed!
ref_out result sample: tensor([-0.0012,  0.0035, -0.0003, -0.0001,  0.0035,  0.0009, -0.0003,  0.0009],
       device='cuda:0', dtype=torch.bfloat16)
flash_mla_out result sample: tensor([-0.0012,  0.0035, -0.0003, -0.0001,  0.0035,  0.0009, -0.0003,  0.0009],
       device='cuda:0')
txl_mla_out result sample: tensor([-0.0012,  0.0035, -0.0003, -0.0001,  0.0035,  0.0009, -0.0003,  0.0009],
       device='cuda:0')
===============================
Running performance test...
FLASH MLA: 0.485 ms, 621 TFLOPS, 1436 GB/s
TXL MLA: 0.498 ms, 605 TFLOPS, 1398 GB/s

mla_prefill log:

================
Running on TestParam(b=1, s_q=64, s_kv=128, topk=128, h_q=128, h_kv=1, d_qk=576, d_v=512, seed=0, check_correctness=True, benchmark=True)
FlashMLA Prefill:    17 us, 135.404 TFlops
TXL MLA Prefill:    122 us, 18.744 TFlops

[120L021326]

Only when Batch % NUM_SMS == 0, can txl version decoding run. On modal the SM count is 132, so I set test Batch is 132.
PCle version H100 has 114 SMs.
TO DO: support varied Batch and varied kv_seq_len

[120L021326]

tl.store block the Batch loop