bugfix: MLA decode should multiply sm_scale by math::log2e (#787)

Author: tsu-bin

flashinfer/jit/attention.py

@@ -180,6 +180,7 @@ def gen_batch_decode_mla_module(
         dtype_o,
         dtype_idx,
         head_dim,
+        head_dim,
         use_sliding_window,
         use_logits_soft_cap,
     )

include/flashinfer/attention/decode.cuh

@@ -857,6 +857,7 @@ __global__ void BatchDecodeWithPagedKVCacheKernelMLA(Params params) {
   const float rope_rcp_scale = params.rope_rcp_scale;
   const float rope_rcp_theta = params.rope_rcp_theta;
   const bool partition_kv = params.partition_kv;
+  params.sm_scale *= math::log2e;
 
   constexpr uint32_t head_dim_ckv = bdx * vec_size_ckv;
   constexpr uint32_t head_dim_kpe = bdx * vec_size_kpe;