[MLA] use manually tuned block size for MLA (#2021)

Signed-off-by: Guangxiang Du <gxd@google.com>

Author: gxd3

tests/layers/common/test_attention_interface.py

@@ -338,11 +338,6 @@ def test_mla_attention(monkeypatch, mesh):
         request_distribution=jnp.array([0, 0, NUM_SEQS], dtype=jnp.int32),
     )
 
-    mock_tuned_block_sizes = MagicMock(return_value=(8, 8))
-    monkeypatch.setattr(
-        "tpu_inference.layers.common.attention_interface.get_tuned_block_sizes",
-        mock_tuned_block_sizes)
-
     expected_output = jnp.full(q_TNA.shape, 0.5)
     expected_new_cache = jnp.full(kv_cache_shape, 0.1)
 
@@ -365,15 +360,13 @@ def test_mla_attention(monkeypatch, mesh):
         sm_scale=0.1,
     )
 
-    # Verify mocked functions were called
-    mock_tuned_block_sizes.assert_called_once()
     mock_mla_kernel.assert_called_once()
 
     # Verify output correctness
     assert jnp.array_equal(output, expected_output)
     assert jnp.array_equal(final_kv_cache, expected_new_cache)
 
     _, kernel_kwargs = mock_mla_kernel.call_args
-    assert kernel_kwargs["num_kv_pages_per_block"] == 4
-    assert kernel_kwargs["num_queries_per_block"] == 4
+    assert kernel_kwargs["num_kv_pages_per_block"] == 3
+    assert kernel_kwargs["num_queries_per_block"] == 1
     assert kernel_kwargs["sm_scale"] == 0.1

tests/platforms/test_tpu_platform.py

@@ -251,6 +251,7 @@ def test_check_and_update_config_block_size(self, mock_logger, mock_update,
         vllm_config.cache_config = MagicMock()
         vllm_config.cache_config.user_specified_block_size = False
         vllm_config.cache_config.block_size = 16
+        vllm_config.model_config.use_mla = False
 
         with patch.dict(
                 'sys.modules', {

tpu_inference/layers/common/attention_interface.py

@@ -31,8 +31,6 @@
 from tpu_inference import envs
 from tpu_inference.kernels.flash_attention.kernel import flash_attention
 from tpu_inference.kernels.mla.v2.kernel import mla_ragged_paged_attention
-from tpu_inference.kernels.ragged_paged_attention.v3.tuned_block_sizes import \
-    get_tuned_block_sizes
 from tpu_inference.layers.common.attention_metadata import AttentionMetadata
 from tpu_inference.layers.common.sharding import ShardingAxisName
 from tpu_inference.logger import init_logger
@@ -521,16 +519,9 @@ def mla_attention(
     )
 
     def _mla_ragged_paged_attention(q, q_rope, k, k_rope, cache, *args):
-        max_num_tokens = q.shape[0]
-        max_num_seqs = md.seq_lens.shape[0]
-        pages_per_seq = md.block_tables.shape[0] // max_num_seqs
-
-        bkv_p, bq_sz = get_tuned_block_sizes(q.dtype, cache.dtype,
-                                             num_attention_heads, 1,
-                                             qk_nope_head_dim, cache.shape[1],
-                                             max_num_tokens, pages_per_seq)
-        num_kv_pages_per_block = min(min(pages_per_seq, bkv_p), 4)
-        num_queries_per_block = min(min(max_num_tokens, bq_sz), 4)
+        # TODO: use auto tuner to find the best block sizes.
+        num_kv_pages_per_block = 3
+        num_queries_per_block = 1
 
         out, new_cache = mla_ragged_paged_attention(
             q,

tpu_inference/layers/vllm/backends/flash_attn_mla.py

@@ -17,6 +17,7 @@
 import torch
 from jax.sharding import Mesh
 from torchax.interop import jax_view
+from vllm.config import VllmConfig
 from vllm.model_executor.layers.attention.mla_attention import MLAAttention
 from vllm.v1.attention.backend import (AttentionBackend, AttentionLayer,
                                        MLAAttentionImpl)
@@ -43,6 +44,10 @@ def get_name() -> str:
     def get_impl_cls() -> type["PallasMLAttentionBackend"]:
         return PallasMLAttentionBackendImpl
 
+    @staticmethod
+    def get_page_size(vllm_config: VllmConfig) -> int:
+        return 1024
+
 
 class PallasMLAttentionBackendImpl(MLAAttentionImpl):
 

tpu_inference/platforms/tpu_platform.py

@@ -212,19 +212,25 @@ def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
         # For v0, the default block size is 16.
         if cache_config and not cache_config.user_specified_block_size:
             if vllm_config.model_config:
-                from tpu_inference.layers.vllm.backends.flash_attn import \
-                    PallasAttentionBackend
-                cache_config.block_size = PallasAttentionBackend.get_page_size(
-                    vllm_config)  # type: ignore[assignment]
-                min_page_size = PallasAttentionBackend.get_min_page_size(
-                    vllm_config)
-                if min_page_size > cache_config.block_size:
-                    logger.warning(
-                        "Increase the page size from %s to %s to avoid SMEM OOM",
-                        cache_config.block_size,
-                        min_page_size,
-                    )
-                    cache_config.block_size = min_page_size  # type: ignore[assignment]
+                if vllm_config.model_config.use_mla:
+                    from tpu_inference.layers.vllm.backends.flash_attn_mla import \
+                        PallasMLAttentionBackend
+                    cache_config.block_size = PallasMLAttentionBackend.get_page_size(
+                        vllm_config)  # type: ignore[assignment]
+                else:
+                    from tpu_inference.layers.vllm.backends.flash_attn import \
+                        PallasAttentionBackend
+                    cache_config.block_size = PallasAttentionBackend.get_page_size(
+                        vllm_config)  # type: ignore[assignment]
+                    min_page_size = PallasAttentionBackend.get_min_page_size(
+                        vllm_config)
+                    if min_page_size > cache_config.block_size:
+                        logger.warning(
+                            "Increase the page size from %s to %s to avoid SMEM OOM",
+                            cache_config.block_size,
+                            min_page_size,
+                        )
+                        cache_config.block_size = min_page_size  # type: ignore[assignment]
             logger.info(
                 f"Using KV cache block size: {cache_config.block_size}")