[KVCache] PagedKVCache refactor, FlashInfer JIT and MLA integration

This PR consists of the following parts:

* We reorganized `paged_kv_cache.cc` by moving some of the utilities
to `attn_utils.h`.

* To integrate with the JIT kernel compilation in the latest FlashInfer
project, while still being able to support attention kernels written
with TIR, we introduced `AttnBackendFunc` in `attn_backend.h`, which
exposes attention interfaces (e.g., `MHA`, `MLA`) to PagedKVCache.
We subclass `AttnBackendFunc` and implement FlashInfer backends and
TIR backends respectively.

* With `AttnBackendFunc`, we refactored the PagedKVCache constructor.
The new constructor is not backward compatible, and will break the
existing compiled model libraries.

* For both TIR and FlashInfer attention implementations, now we
require an explicit attention softmax scale factor `sm_scale` to be
passed in. Previously, it has an inlined `sm_scale` of
`head_dim ** -0.5`. Due to the recent LLM inference techniques such
as MLA weight absorption in DeepSeek models, the inlined `sm_scale`
causes confusion and inconvenience. To keep attention interface
standard and clear, we now require the explicit passing of `sm_scale`.

* We refactored the existing GPU unit tests of the PagedKVCache,
by updating from numpy to PyTorch for std calculation. This
significantly reduces the test case run time.

Author: MasterJH5574

docs/how_to/tutorials/optimize_llm.py

@@ -191,7 +191,9 @@ def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id:
         qkv = op.reshape(qkv, (b, s, h_q + h_kv + h_kv, d))
         # Attention
         output = op.reshape(
-            paged_kv_cache.attention_with_fused_qkv(layer_id, qkv, self.num_q_heads),
+            paged_kv_cache.attention_with_fused_qkv(
+                layer_id, qkv, self.num_q_heads, sm_scale=self.head_dim**-0.5
+            ),
             (b, s, h_q * d),
         )
         # Output Projection
@@ -285,6 +287,7 @@ def create_tir_paged_kv_cache(
         page_size: tir.Var,
     ) -> PagedKVCache:
         return TIRPagedKVCache(
+            attn_kind="mha",
             max_batch_size=max_batch_size,
             max_total_seq_len=max_total_seq_len,
             prefill_chunk_size=prefill_chunk_size,
@@ -294,7 +297,10 @@ def create_tir_paged_kv_cache(
             num_hidden_layers=self.num_hidden_layers,
             num_attention_heads=self.num_attention_heads,
             num_key_value_heads=self.num_key_value_heads,
-            head_dim=self.head_dim,
+            qk_head_dim=self.head_dim,
+            v_head_dim=self.head_dim,
+            mla_original_qk_head_dim=0,
+            mla_original_v_head_dim=0,
             rope_mode=RopeMode.NORMAL,
             rope_scale=1,
             rope_theta=self.rope_theta,

python/tvm/relax/__init__.py

@@ -91,6 +91,7 @@
 )
 
 # pipeline
+from .pipeline import get_default_pipeline
 from .pipeline import get_pipeline
 from .pipeline import register_pipeline
 

python/tvm/relax/backend/cuda/__init__.py

@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """The Relax CUDA backend compilation pipeline and other passes."""
+from . import flashinfer
 from .pipeline import (
     finalize_passes,
     get_default_pipeline,