deepseek sharding and mla attention plumbing

khatwanimohit · khatwanimohit · commit e09523deb76d · 2026-03-11T17:46:40.000Z
diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml
@@ -463,6 +463,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']],
@@ -1110,6 +1111,8 @@ vllm_hf_config_path: ""
 # A JSON string of overrides to apply to the HuggingFace-style config for the vLLM adapter.
 # This can be used to override specific settings without modifying the original config file. 
 vllm_hf_overrides: {}
+# 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]
diff --git a/src/maxtext/configs/inference/vllm_deepseek.yml b/src/maxtext/configs/inference/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']],
+                    ]
diff --git a/src/maxtext/configs/post_train/rl.yml b/src/maxtext/configs/post_train/rl.yml
@@ -154,6 +154,8 @@ max_num_seqs: null
 async_scheduling: True
 # stop generation when any of these strings is generated
 stop_strings: null
+# path to initialize vllm config
+vllm_config_path: 'src/maxtext/configs/inference/vllm.yml'
 
 # ====== Checkpoint Configuration ======
 enable_checkpointing: True
diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py
@@ -1610,6 +1610,7 @@ class VLLM(BaseModel):
       description="Overrides for HuggingFace model config for MaxText model.",
   )
   vllm_hf_config_path: str = Field("", description="Path to HuggingFace model config for MaxText model.")
+  vllm_config_path: str = Field("src/maxtext/configs/inference/vllm.yml", description="path to yaml file for loading vLLM config.")
 
 
 class RL(BaseModel):
diff --git a/src/maxtext/inference/vllm_decode.py b/src/maxtext/inference/vllm_decode.py
@@ -29,6 +29,7 @@
       use_chat_template=True
 """
 
+import copy
 import os
 from typing import Any, Sequence
 
@@ -67,6 +68,21 @@ def decode_with_vllm(config: Config) -> None:
     config: MaxText config.
   """
   # Prepare vLLM Arguments
+  # Use user-provided vllm_additional_config as base (includes model-specific
+  # overrides like base_num_decoder_layers, override_model_config, etc.), then
+  # fill in defaults and runtime-derived values on top.
+  additional_config = copy.deepcopy(config.vllm_additional_config) if config.vllm_additional_config else {}
+  additional_config.setdefault("maxtext_config", {})
+  additional_config["maxtext_config"].setdefault("model_name", config.model_name)
+  additional_config["maxtext_config"].setdefault("weight_dtype", "bfloat16")
+  additional_config["maxtext_config"].setdefault("allow_split_physical_axes", True)
+  additional_config["maxtext_config"]["debug_sharding"] = config.debug_sharding
+  additional_config.setdefault("sharding", {})
+  additional_config["sharding"].setdefault("sharding_strategy", {})
+  additional_config["sharding"]["sharding_strategy"].setdefault("enable_dp_attention", config.enable_dp_attention)
+  # Pass vllm_config_path so the adapter can use it as the MaxText base config.
+  additional_config.setdefault("vllm_config_path", str(config.vllm_config_path))
+
   vllm_args = {
       "model": config.tokenizer_path,
       "max_model_len": config.max_target_length,
@@ -76,19 +92,7 @@ def decode_with_vllm(config: Config) -> None:
       "hf_overrides": config.vllm_hf_overrides,
       "gpu_memory_utilization": config.hbm_utilization_vllm,
       "async_scheduling": config.async_scheduling,
-      "additional_config": {
-          "maxtext_config": {
-              "model_name": config.model_name,
-              "weight_dtype": "bfloat16",
-              "allow_split_physical_axes": True,
-              "debug_sharding": config.debug_sharding,
-          },
-          "sharding": {
-              "sharding_strategy": {
-                  "enable_dp_attention": config.enable_dp_attention,
-              },
-          },
-      },
+      "additional_config": additional_config,
   }
 
   if config.load_parameters_path:
@@ -106,8 +110,7 @@ def decode_with_vllm(config: Config) -> None:
       f"and EP={config.ici_expert_parallelism if enable_expert_parallel else 1}..."
   )
 
-  vllm_config_path = os.path.join(MAXTEXT_CONFIGS_DIR, "inference", "vllm.yml")
-  argv_list = ["", str(vllm_config_path), "log_config=False"]
+  argv_list = ["", str(config.vllm_config_path), "log_config=False"]
   vllm_config = pyconfig.initialize(argv_list)
 
   with nn_partitioning.axis_rules(vllm_config.logical_axis_rules):
@@ -145,7 +148,7 @@ def decode_with_vllm(config: Config) -> None:
       max_tokens=max_tokens_to_generate,
       top_k=config.decode_sampling_top_k,
       top_p=config.decode_sampling_nucleus_p,
-      seed=FLAGS.seed,
+      # seed=FLAGS.seed,
   )
 
   outputs = llm.generate(prompts, sampling_params)
diff --git a/src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py b/src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py
@@ -69,8 +69,12 @@ def generate_maxtext_config(vllm_config: VllmConfig) -> pyconfig.HyperParameters
       )
       overrides["load_parameters_path"] = None
 
-  # Add base config path to positional args
-  base_config_path = os.path.join(MAXTEXT_CONFIGS_DIR, "inference", "vllm.yml")
+  # Add base config path to positional args — prefer the caller-supplied
+  # vllm_config_path from additional_config, fall back to vllm.yml default.
+  base_config_path = vllm_config.additional_config.get(
+      "vllm_config_path",
+      os.path.join(MAXTEXT_CONFIGS_DIR, "inference", "vllm.yml"),
+  )
   argv_list = ["", str(base_config_path)]
 
   maxtext_config = pyconfig.initialize(argv_list, **overrides)
diff --git a/src/maxtext/layers/attention_mla.py b/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
 
@@ -623,7 +625,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."""
@@ -941,15 +947,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,
@@ -1005,7 +1114,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")
@@ -1039,7 +1148,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:
+      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,
diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py
diff --git a/src/maxtext/models/deepseek.py b/src/maxtext/models/deepseek.py
diff --git a/src/maxtext/trainers/post_train/rl/train_rl.py b/src/maxtext/trainers/post_train/rl/train_rl.py

Original file line number	Diff line number	Diff line change
`@@ -1610,6 +1610,7 @@ class VLLM(BaseModel):`
`1610`	`1610`	`description="Overrides for HuggingFace model config for MaxText model.",`
`1611`	`1611`	`)`
`1612`	`1612`	`vllm_hf_config_path: str = Field("", description="Path to HuggingFace model config for MaxText model.")`
	`1613`	`+ vllm_config_path: str = Field("src/maxtext/configs/inference/vllm.yml", description="path to yaml file for loading vLLM config.")`
`1613`	`1614`
`1614`	`1615`
`1615`	`1616`	`class RL(BaseModel):`