deepseek sharding and mla attention plumbing

Author: khatwanimohit

src/MaxText/__init__.py

@@ -34,7 +34,6 @@
 from maxtext.trainers.post_train.dpo import dpo_utils
 from maxtext.utils import maxtext_utils
 from maxtext.utils import model_creation_utils
-from maxtext.utils.model_creation_utils import from_config
 
 Transformer = models.Transformer
 transformer_as_linen = models.transformer_as_linen

src/MaxText/configs/base.yml

@@ -430,6 +430,7 @@ logical_axis_rules: [
                       ['decode_length', ['sequence']],
                       ['mlp', ['fsdp_transpose', 'tensor', 'tensor_sequence', 'autoregressive']],
                       ['mlp_no_fsdp', ['tensor', 'tensor_sequence', 'autoregressive']],
+                      ['moe_mlp', ['fsdp_transpose', 'tensor', 'tensor_sequence', 'autoregressive']],
                       ['vocab', ['tensor', 'tensor_transpose', 'tensor_sequence', 'autoregressive']],
                       ['heads', ['tensor', 'tensor_transpose', 'tensor_sequence', 'autoregressive']],
                       ['q_heads', ['tensor', 'tensor_transpose', 'tensor_sequence', 'autoregressive']],
@@ -1047,6 +1048,8 @@ use_jax_splash: false
 # vLLM Adapter Configurations
 # Path to the HuggingFace-style config directory for the adapter (e.g. src/MaxText/integration/vllm/maxtext_vllm_adapter)
 vllm_hf_config_path: ""
+# Path to yaml file for loading vLLM config
+vllm_config_path: ""
 # JSON string containing additional configuration for the vLLM model (e.g. '{"maxtext_config": {...}}')
 vllm_additional_config: {}
 # When use_jax_splash=True, force the layout of the query tensor to be [..., NUM_HEADS, HEAD_DIM, SEQ_LENGTH]

src/MaxText/configs/rl.yml

@@ -149,6 +149,8 @@ enable_dp_attention: False
 # Performance tuning for samplers
 max_num_batched_tokens: null
 max_num_seqs: null
+# path to initialize vllm config
+vllm_config_path: 'src/MaxText/configs/vllm.yml'
 
 # ====== Checkpoint Configuration ======
 enable_checkpointing: True

src/MaxText/configs/types.py

@@ -1490,6 +1490,7 @@ class VLLM(BaseModel):
   max_num_seqs: Optional[int] = Field(None, description="Max number of sequences in vLLM.")
   vllm_additional_config: dict[str, Any] = Field(default_factory=dict, description="Additional vLLM config options.")
   vllm_hf_config_path: str = Field("", description="Path to HuggingFace model config for MaxText model.")
+  vllm_config_path: str = Field("src/MaxText/configs/vllm.yml", description="path to yaml file for loading vLLM config.")
 
 
 class RL(BaseModel):

src/MaxText/configs/vllm.yml

@@ -53,6 +53,7 @@ logical_axis_rules: [
                       ['decode_length', []],
                       ['mlp', ['model', 'attn_dp']],
                       ['mlp_no_fsdp', ['model', 'attn_dp']],
+                      ['moe_mlp', ['model', 'attn_dp']],
                       ['vocab', ['model', 'attn_dp']],
                       ['heads', ['model']],
                       ['q_heads', ['model']],

src/MaxText/configs/vllm_deepseek.yml

@@ -0,0 +1,69 @@
+# Copyright 2023–2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+base_config: "vllm.yml"
+
+logical_axis_rules: [
+                      ['activation_batch', ['']],
+                      ['activation_batch_no_exp', []],
+                      ['activation_embed_and_logits_batch', ['expert']],
+                      ['activation_embed_and_logits_batch_sequence', ['expert']],
+                      ['activation_heads', ['model']],
+                      ['activation_kv_heads', ['model']],
+                      ['activation_attn_length', ['expert']],
+                      ['activation_attn_length_no_exp', []],
+                      ['activation_length', ['data', 'expert']],
+                      ['activation_length_no_exp', 'data'],
+                      ['activation_q_length', ['expert']],
+                      ['activation_attn_embed', 'model'],
+                      ['activation_embed', ['model', 'attn_dp']],
+                      ['activation_mlp', ['model', 'attn_dp', 'expert']],
+                      ['activation_kv', ['model']],
+                      ['activation_prefill_kv_batch', ['expert']],
+                      ['activation_kv_batch', ['']],
+                      ['activation_kv_batch_no_exp', []],
+                      ['activation_kv_head_dim', ['model', 'attn_dp', 'expert']],
+                      ['activation_vocab', ['model', 'attn_dp']],
+                      ['activation_norm_length', []],
+                      ['activation_exp', ['expert']],
+                      ['decode_batch', ['expert']],
+                      ['decode_length', []],
+                      ['mlp_no_fsdp', ['model', 'attn_dp', 'expert']],
+                      ['vocab', ['model', 'attn_dp', 'expert']],
+                      ['heads', ['expert', 'attn_dp', 'model']],
+                      ['q_heads', []],
+                      ['kv_heads', []],
+                      ['kv_head_dim', ['model', 'attn_dp', 'expert']],
+                      ['kv', ['model', 'attn_dp', 'expert']],
+                      ['kv', []],
+                      ['embed', []],
+                      ['mlp', ['model', 'attn_dp', 'expert']],
+                      ['moe_mlp', []],
+                      ['embed_tensor_transpose', ['attn_dp', 'model']],
+                      ['embed_no_exp', []],
+                      ['q_lora', []],
+                      ['kv_lora', []],
+                      ['norm', []],
+                      ['cache_heads', ['model']],
+                      ['exp', ['expert', 'attn_dp', 'model']],
+                      ['paged_kv_heads', ['model']],
+                      ['cache_batch_prefill', []],
+                      ['cache_batch', []],
+                      ['cache_sequence', []],
+                      ['cache_heads_none', []],
+                      ['cache_kv', []],
+                      ['kv_lora_up_proj',['expert', 'attn_dp', 'model']],
+                      ['q_lora_up_proj',['expert', 'attn_dp', 'model']],
+                    ]

src/MaxText/kernels/sort_activations.py

@@ -73,9 +73,7 @@ def _unroute_fwd(
   )
 
 
-def _unroute_bwd(
-    use_custom_mosaic_kernel: bool, residuals: jax.Array, grads: jax.Array
-) -> tuple[jax.Array, None]:
+def _unroute_bwd(use_custom_mosaic_kernel: bool, residuals: jax.Array, grads: jax.Array) -> tuple[jax.Array, None]:
   selected_experts = residuals
   return _route_impl(grads, selected_experts, use_custom_mosaic_kernel), None
 
@@ -90,8 +88,7 @@ def _route_impl(
 ) -> jax.Array:
   """Gather `tokens` according to `selected_experts`."""
   assert (
-      tokens.shape[0] == selected_experts.shape[0]
-      and selected_experts.ndim == 2
+      tokens.shape[0] == selected_experts.shape[0] and selected_experts.ndim == 2
   ), f"{tokens.shape=}, {selected_experts.shape=}"
   if use_custom_mosaic_kernel:
     raise NotImplementedError("Custom Mosaic kernel not implemented.")
@@ -104,10 +101,8 @@ def _unroute_impl(
     selected_experts: jax.Array,
     use_custom_mosaic_kernel: bool,
 ) -> jax.Array:
-  assert (
-      tokens.shape[0] == selected_experts.shape[0] * selected_experts.shape[1]
-      and selected_experts.ndim == 2
-  )
+  """Reverse the routing operation, restoring tokens to their original order."""
+  assert tokens.shape[0] == selected_experts.shape[0] * selected_experts.shape[1] and selected_experts.ndim == 2
   inds = jnp.argsort(jnp.argsort(jnp.ravel(selected_experts)))
   return jnp.sum(
       jnp.reshape(
@@ -118,9 +113,7 @@ def _unroute_impl(
   )
 
 
-def _sort_impl(
-    tokens: jax.Array, inds: jax.Array, use_custom_mosaic_kernel: bool
-) -> jax.Array:
+def _sort_impl(tokens: jax.Array, inds: jax.Array, use_custom_mosaic_kernel: bool) -> jax.Array:
   if use_custom_mosaic_kernel:
     raise NotImplementedError("Custom Mosaic kernel not implemented.")
   else:

src/MaxText/layers/attention_mla.py

@@ -21,6 +21,8 @@
 import jax
 from jax.ad_checkpoint import checkpoint_name
 from jax.experimental import layout
+from jax.sharding import PartitionSpec as P
+from jax.experimental import shard_map
 import jax.numpy as jnp
 from jax.sharding import Mesh, NamedSharding
 
@@ -619,7 +621,11 @@ def __init__(
       )
 
     # Module attribute names must match names previously passed to Linen for checkpointing
-    self.MlaKVCache_0 = self.init_mla_kv_caches(inputs_kv_shape) if model_mode != MODEL_MODE_TRAIN else None
+    self.MlaKVCache_0 = (
+        self.init_mla_kv_caches(inputs_kv_shape)
+        if model_mode != MODEL_MODE_TRAIN and config.attention != "vllm_rpa"
+        else None
+    )
 
   def _init_projections(self, inputs_q_shape: Tuple, inputs_kv_shape: Tuple) -> None:
     """Initializes the MLA-specific projections."""
@@ -937,15 +943,118 @@ def mla_kv_projection(self, inputs: Array, inputs_positions: Array, decoder_segm
 
     key, value = self.mla_get_key_value(low_rank_main, key_rope, model_mode)
     cached_values = [None, None]
-    if self.config.attention != "paged" and model_mode != MODEL_MODE_TRAIN:
+    if self.config.attention != "paged" and self.config.attention != "vllm_rpa" and model_mode != MODEL_MODE_TRAIN:
       if self.config.mla_naive_kvcache:
         cached_values = self.update_kv_caches(key, value, decoder_segment_ids, model_mode, previous_chunk)
       else:
         cached_values = self.update_mla_kv_caches(
             low_rank_main, key_rope, decoder_segment_ids, model_mode, previous_chunk
         )
 
-    return key, value, cached_values
+    return key, value, cached_values, low_rank_main, key_rope
+
+  def mla_rpa_vllm(self, q_nope, q_rope, k_latent, k_rope, mla_kv_cache, mla_metadata):
+    """Forward function for vLLM serving with MLA attention.
+
+    Args:
+        q_nope: Query nope part [T, N, qk_nope_head_dim]
+        q_rope: Query rope part [T, N, qk_rope_head_dim]
+        k_latent: Latent KV representation [S, kv_lora_rank] (NOT expanded k_nope)
+        k_rope: Key rope part [S, qk_rope_head_dim] (NO head dimension)
+        mla_kv_cache: The KV cache
+        mla_metadata: Attention metadata
+    """
+    md = mla_metadata
+    try:
+      # pylint: disable=import-outside-toplevel
+      # pytype: disable=import-error
+      from tpu_inference.kernels.mla.v1.kernel import mla_ragged_paged_attention
+      from tpu_inference.kernels.ragged_paged_attention.v3.tuned_block_sizes import get_tuned_block_sizes
+    except ImportError as e:
+      raise ImportError(
+          "vLLM RPA attention ops require the vllm-tpu package. Please install it with `pip install vllm-tpu`."
+      ) from e
+
+    if mla_kv_cache is None or mla_metadata is None:
+      raise ValueError("kv_cache and attention_metadata must be provided when using vLLM.")
+
+    wkv_b_kernel = self.wkv_b.kernel.value
+    wk_b_kernel = wkv_b_kernel[..., : self.qk_nope_head_dim]
+    wv_b_kernel = wkv_b_kernel[..., self.qk_nope_head_dim :]
+    q_absorbed = jnp.einsum("TNH,ANH->TNA", q_nope, wk_b_kernel)
+
+    def _mla_ragged_paged_attention(q, q_rope, k, k_rope, kv_cache, *args):
+      def _initialize_block_sizes():
+        # Set reasonable starting estimates for block sizes. (TODO(gpolovets): update this to use tuned sizes)
+        max_num_tokens = q_absorbed.shape[0]
+        max_num_seqs = md.seq_lens.shape[0]
+        num_page_indices = md.block_tables.shape[0]
+        assert num_page_indices % max_num_seqs == 0
+        pages_per_seq = num_page_indices // max_num_seqs
+        # num_kv_pages_per_block = min(pages_per_seq, 16)
+        bkv_p, bq_sz = get_tuned_block_sizes(
+            q_nope.dtype,
+            q_nope.dtype,  # changed to q_nope dtype from mla_kv_cache.dtype
+            self.num_query_heads,
+            1,  # num_kv_heads for MLA kernel
+            self.qk_nope_head_dim,
+            q_nope.shape[1],  # page size ?? kv_cache.shape[1]
+            max_num_tokens,
+            pages_per_seq,
+        )
+        num_kv_pages_per_block = min(pages_per_seq, bkv_p, 4)
+        num_queries_per_block = min(max_num_tokens, bq_sz, 4)  # OOMS at 8
+        return num_kv_pages_per_block, num_queries_per_block
+
+      num_kv_pages_per_block, num_queries_per_block = _initialize_block_sizes()
+      output, kv_cache = mla_ragged_paged_attention(
+          q,
+          q_rope,
+          k,
+          k_rope,
+          kv_cache,
+          *args,
+          sm_scale=1.0,
+          num_kv_pages_per_block=num_kv_pages_per_block,
+          num_queries_per_block=num_queries_per_block,
+      )
+      return kv_cache, output
+
+    in_specs = (
+        P(("attn_dp", "model", "expert"), None, None),  # q
+        P(("attn_dp", "model", "expert"), None, None),  # q_rope
+        P(("attn_dp", "model", "expert"), None),  # k
+        P(("attn_dp", "model", "expert"), None),  # k_rope
+        P(("attn_dp", "model", "expert")),  # kv_cache
+        P(("data", "attn_dp")),  # md.seq_lens: Replicated
+        P(("data", "attn_dp")),  # page_indices_flat: Replicated
+        P(("data", "attn_dp")),  # query_start_loc: Replicated
+        P(("data", "attn_dp")),  # distribution: Replicated
+    )
+
+    out_specs = (P(("attn_dp", "model", "expert"), None, None), P(("attn_dp", "model", "expert")))
+
+    kv_cache, output = jax.jit(
+        shard_map.shard_map(
+            _mla_ragged_paged_attention,
+            mesh=self.mesh,
+            in_specs=in_specs,
+            out_specs=out_specs,
+            check_rep=False,
+        ),
+    )(
+        q_absorbed,
+        q_rope,
+        k_latent,
+        k_rope,
+        mla_kv_cache,
+        md.seq_lens,
+        md.block_tables,
+        md.query_start_loc,
+        md.request_distribution,
+    )
+    output = jnp.einsum("TNA,ANH->TNH", output, wv_b_kernel)
+    return kv_cache, output
 
   def __call__(
       self,
@@ -1001,7 +1110,7 @@ def __call__(
     query, low_rank_q = self.mla_query_projection(inputs_q, inputs_positions, model_mode)
     if self.config.force_q_layout:
       query = layout.with_layout_constraint(query, DLL(major_to_minor=(0, 2, 3, 1)))
-    key, value, cached_values = self.mla_kv_projection(
+    key, value, cached_values, low_rank_main, key_rope = self.mla_kv_projection(
         inputs_kv, inputs_positions, decoder_segment_ids, model_mode, previous_chunk
     )
     query = checkpoint_name(query, "query_proj")
@@ -1034,8 +1143,22 @@ def __call__(
       )
       unnormalized_out = unnormalized_out[..., : self.v_head_dim]
       out = unnormalized_out / (exp_sum + 1e-9) if exp_sum is not None else unnormalized_out
+    elif self.config.attention == "vllm_rpa" and model_mode != MODEL_MODE_TRAIN and kv_cache is not None:
+      batch, seq_len, num_heads, _ = query.shape
+      query = query.reshape(-1, query.shape[2], query.shape[3])
+      q_nope, q_rope = jnp.split(query, [self.qk_nope_head_dim], axis=-1)
+
+      k_latent = low_rank_main.reshape(-1, self.kv_lora_rank)
+      k_rope_squeezed = key_rope.reshape(-1, self.qk_rope_head_dim)
+
+      updated_kv, attn_out = self.mla_rpa_vllm(
+          q_nope, q_rope, k_latent, k_rope_squeezed, mla_kv_cache=kv_cache, mla_metadata=attention_metadata
+      )
+      out = attn_out.reshape(batch, seq_len, num_heads, self.v_head_dim)
+      kv_cache = updated_kv
     else:
-      # Pass the index_mask to the Attention Op
+      if self.config.attention == "vllm_rpa" and kv_cache is None and model_mode != MODEL_MODE_TRAIN:
+        model_mode = MODEL_MODE_TRAIN
       out = self.attention_op(query, key, value, decoder_segment_ids, model_mode, cached_values, index_mask=index_mask)
 
     if model_mode == MODEL_MODE_TRAIN and self.config.expert_shard_attention_option == EP_AS_CONTEXT: