batched rpa kernel

Author: khatwanimohit

src/maxtext/configs/pyconfig_deprecated.py

@@ -106,6 +106,7 @@ def validate_attention_kernel(s: str) -> None:
       "cudnn_flash_te",
       "cudnn_flash_jax",
       "vllm_rpa",
+      "vllm_batched_rpa",
   )
   if s not in valid_attention_kernels:  # currently supported attention
     raise ValueError("Invalid attention kernel was passed. Valid options ", valid_attention_kernels)

src/maxtext/configs/types.py

@@ -561,7 +561,7 @@ class Attention(BaseModel):
 
   attention: str = Field(
       "autoselected",
-      description="The attention algorithm to use (dot_product, flash, etc).",
+      description="The attention algorithm to use (dot_product, flash, cudnn_flash_te, vllm_rpa, vllm_batched_rpa, etc).",
   )
   attention_type: Literal["global", "local_sliding", "chunk", "mla", "full", "compressed"] = Field(
       "global", description="The variant of attention to use."

src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py

@@ -90,6 +90,9 @@ def generate_maxtext_config(vllm_config: VllmConfig) -> pyconfig.HyperParameters
       )
       overrides["load_parameters_path"] = None
 
+  if overrides.get("attention") == "vllm_batched_rpa" or overrides.get("use_batched_rpa", False):
+    os.environ["USE_BATCHED_RPA_KERNEL"] = "1"
+
   # Add base config path to positional args
   base_config_path = os.path.join(MAXTEXT_CONFIGS_DIR, "inference", "vllm.yml")
   argv_list = ["", str(base_config_path)]

src/maxtext/layers/attention_op.py

@@ -481,7 +481,14 @@ def __init__(
     self.kv_quant = kv_quant
     self.attention_type = attention_type
     # Block sizes are only used by TPU splash attention kernels. Exclude non-splash kernels
-    if self.attention_kernel not in ("dot_product", "paged", "vllm_rpa", "cudnn_flash_te", "cudnn_flash_jax"):
+    if self.attention_kernel not in (
+        "dot_product",
+        "paged",
+        "vllm_rpa",
+        "vllm_batched_rpa",
+        "cudnn_flash_te",
+        "cudnn_flash_jax",
+    ):
       if self.attention_type == AttentionType.LOCAL_SLIDING:
         self.block_q = self.config.local_sa_block_q
         self.block_kv = self.config.local_sa_block_kv
@@ -1069,7 +1076,7 @@ def apply_attention(
         or (self.attention_kernel == "autoselected" and model_mode == MODEL_MODE_AUTOREGRESSIVE)
         or (self.attention_kernel == "autoselected" and length < 128)
         or (self.attention_kernel == "paged")
-        or (self.attention_kernel == "vllm_rpa")
+        or (self.attention_kernel in ("vllm_rpa", "vllm_batched_rpa"))
     ):
       return self.apply_attention_dot(
           query,

src/maxtext/layers/attentions.py

@@ -16,6 +16,7 @@
 
 import dataclasses
 import functools
+import os
 from typing import Any, Iterable, Optional, Tuple, Union, cast
 
 from jax.ad_checkpoint import checkpoint_name
@@ -436,7 +437,9 @@ def __init__(
     # Module attribute names must match names previously passed to Linen for checkpointing
     self.KVCache_0 = (
         self.init_kv_caches(inputs_kv_shape=inputs_kv_shape)
-        if self.model_mode != MODEL_MODE_TRAIN and base_kv_cache and config.attention != "vllm_rpa"
+        if self.model_mode != MODEL_MODE_TRAIN
+        and base_kv_cache
+        and config.attention not in ("vllm_rpa", "vllm_batched_rpa")
         else None
     )
 
@@ -1014,6 +1017,8 @@ def forward_serve_vllm(
       rpa_metadata: dict[str, Any] | None = None,
   ) -> tuple[Array, list[Array]]:
     """Forward function for vLLM serving with RPA attention."""
+    if self.config.attention == "vllm_batched_rpa":
+      os.environ["USE_BATCHED_RPA_KERNEL"] = "1"
     try:
       # pylint: disable=import-outside-toplevel
       # pytype: disable=import-error
@@ -1214,7 +1219,7 @@ def __call__(
 
     assert not self.config.quantize_kvcache or self.kv_quant
 
-    if self.config.attention == "vllm_rpa" and model_mode != MODEL_MODE_TRAIN:
+    if self.config.attention in ("vllm_rpa", "vllm_batched_rpa") and model_mode != MODEL_MODE_TRAIN:
       batch, seq_len, num_heads, head_dim = query.shape
       attn_out, updated_kv = self.forward_serve_vllm(
           query, key, value, rpa_kv_cache=kv_cache, rpa_metadata=attention_metadata

src/maxtext/layers/decoders.py

@@ -1269,11 +1269,11 @@ def __call__(
     # When initializing with vLLM RPA attention, we need to run the output head to
     # initialize any parameters associated with it.
     # Same case applicable to vocab tiling
-    if self.is_initializing() and (cfg.num_vocab_tiling > 1 or cfg.attention == "vllm_rpa"):
+    if self.is_initializing() and (cfg.num_vocab_tiling > 1 or cfg.attention in ("vllm_rpa", "vllm_batched_rpa")):
       _ = self.apply_output_head(shared_embedding, hidden_state, deterministic, model_mode)
 
     # When invoking from vLLM with RPA attention, logit computation is deferred to a later stage.
-    if cfg.attention == "vllm_rpa":
+    if cfg.attention in ("vllm_rpa", "vllm_batched_rpa"):
       logits = None
     # When in the Indexer Dense Warm-up stage, skip the expensive output head projection
     # for efficiency, as the main model is frozen and the LM loss is not needed.

src/maxtext/layers/moe.py

@@ -453,13 +453,13 @@ def __init__(
       self.wi_kernel_axes = ("exp", "embed_moe", "mlp_moe")
       self.wo_kernel_axes = ("exp", "mlp_moe", "embed_moe")
 
-    if self.config.attention == "vllm_rpa":
+    if self.config.attention in ("vllm_rpa", "vllm_batched_rpa"):
       # vLLM uses 'model' as the tensor parallelism axis name
       self._tensor_parallelism_name = ("model", "attn_dp")
     else:
       self._tensor_parallelism_name = "tensor"
 
-    if self.config.attention == "vllm_rpa" and self.config.enable_dp_attention:
+    if self.config.attention in ("vllm_rpa", "vllm_batched_rpa") and self.config.enable_dp_attention:
       self._expert_parallelism_name = "attn_dp_expert"
     elif self.config.custom_mesh_and_rule == ctypes.CustomRule.CP_AS_EP:
       # when custom mesh and rule is cp-as-ep, context axis is same with expert in MoE component
@@ -481,7 +481,7 @@ def __init__(
         # tpu-inference applies the score function in the fused_moe_gmm kernel,
         # so we don't apply it here to avoid redundant computation.
         # See https://github.com/vllm-project/tpu-inference/blob/main/tpu_inference/layers/common/fused_moe_gmm.py#L58.
-        score_func="" if self.config.attention == "vllm_rpa" else self.config.routed_score_func,
+        score_func="" if self.config.attention in ("vllm_rpa", "vllm_batched_rpa") else self.config.routed_score_func,
         matmul_precision=self.config.matmul_precision,
         shard_mode=config.shard_mode,
         rngs=self.rngs,
@@ -1426,7 +1426,7 @@ def jax_ragged_dot_gmm(inputs, kernel, tiling, group_sizes, expert_assignments,
     def get_tokamax_group_sizes(group_sizes, inputs, _kernel):
       if self.config.use_qwix_quantization:
         return group_sizes
-      elif self.config.attention == "vllm_rpa":
+      elif self.config.attention in ("vllm_rpa", "vllm_batched_rpa"):
         return group_sizes
       else:
         return tokamax.RaggedDotGroupSizes(
@@ -3066,7 +3066,7 @@ def __call__(
     fused_kernel = None
     w0_kernel = None
     w1_kernel = None
-    if cfg.prefuse_moe_weights and cfg.attention == "vllm_rpa" and not self.is_hash_routing:
+    if cfg.prefuse_moe_weights and cfg.attention in ("vllm_rpa", "vllm_batched_rpa") and not self.is_hash_routing:
       fused_kernel = jnp.asarray(self.wi[...], self.dtype)
     elif cfg.prefuse_moe_weights:
       wi = jnp.asarray(self.wi[...], self.dtype)
@@ -3096,7 +3096,7 @@ def __call__(
     # The fused MoE kernel currently only supports standard Top-K routing with associated
     # weights. Hash routed layers bypass this kernel and fall back
     # to the sparse matmul implementation.
-    if cfg.attention == "vllm_rpa" and not self.is_hash_routing:
+    if cfg.attention in ("vllm_rpa", "vllm_batched_rpa") and not self.is_hash_routing:
       output, lb_loss, bias_updates = self.fused_moe_matmul(
           inputs,
           gate_logits,

src/maxtext/layers/nnx_decoders.py

@@ -1833,7 +1833,7 @@ def pure_layer_fn(graphdef, state_in, y_in, kv_in):
       hidden_state = y
 
     # When invoking from vLLM with RPA attention, logit computation is deferred to a later stage.
-    if cfg.attention == "vllm_rpa":
+    if cfg.attention in ("vllm_rpa", "vllm_batched_rpa"):
       logits = None
 
     # When in the Indexer Dense Warm-up stage, skip the expensive output head projection

src/maxtext/models/models.py

@@ -251,7 +251,7 @@ def __call__(
           model_mode=model_mode,
       )
 
-    if self.config.attention == "vllm_rpa":
+    if self.config.attention in ("vllm_rpa", "vllm_batched_rpa"):
       # In vLLM, logits are computed separately after updating the KV cache.
       return hidden_state, kv_caches
 
@@ -370,7 +370,7 @@ def __init__(
     dummy_decoder_input_tokens = jnp.ones((batch_size, seq_len), dtype=jnp.int32)
     dummy_decoder_positions = jnp.ones((batch_size, seq_len), dtype=jnp.int32)
 
-    if self.config.attention == "vllm_rpa":
+    if self.config.attention in ("vllm_rpa", "vllm_batched_rpa"):
       try:
         # pylint: disable=import-outside-toplevel
         from tpu_inference.layers.common.attention_metadata import AttentionMetadata  # pytype: disable=import-error
@@ -608,7 +608,7 @@ def __call__(
           model_mode=model_mode,
       )
 
-    if self.config.attention == "vllm_rpa":
+    if self.config.attention in ("vllm_rpa", "vllm_batched_rpa"):
       # In vLLM, logits are computed separately after updating the KV cache.
       return hidden_state, kv_caches
 

tests/inference/benchmark_vllm.sh

@@ -0,0 +1,92 @@
+#!/bin/bash
+# Copyright 2026 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
+#
+#     http://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.
+
+# Generic benchmark script for MaxText models using vLLM on TPU.
+#
+# Usage:
+#   bash tests/inference/benchmark_vllm.sh <MODEL_NAME> [ATTENTION] [TP_SIZE]
+#
+# Arguments:
+#   $1 (MODEL_NAME) : MaxText model name (e.g., "qwen3-0.6b", "llama3.1-8b") [Required]
+#   $2 (ATTENTION)  : Attention kernel to benchmark (default: "vllm_rpa")
+#   $3 (TP_SIZE)    : Tensor parallelism size (default: 4)
+#
+# Environment Variable Overrides:
+#   TOKENIZER      : HuggingFace tokenizer/model path (default: "Qwen/Qwen3-0.6B")
+#
+# Examples:
+#   bash tests/inference/benchmark_vllm.sh qwen3-0.6b
+#   bash tests/inference/benchmark_vllm.sh qwen3-0.6b vllm_batched_rpa
+#   TOKENIZER=meta-llama/Llama-3.1-8B bash tests/inference/benchmark_vllm.sh llama3.1-8b vllm_batched_rpa 4
+
+set -e
+
+if [ -z "$1" ]; then
+  echo "Error: MODEL_NAME is required as the first argument."
+  echo "Usage: bash tests/inference/benchmark_vllm.sh <MODEL_NAME> [ATTENTION] [TP_SIZE]"
+  echo "Example: bash tests/inference/benchmark_vllm.sh qwen3-0.6b"
+  exit 1
+fi
+
+MODEL_NAME=$1
+ATTENTION=${2:-${ATTENTION:-"vllm_rpa"}}
+TP_SIZE=${3:-${TP_SIZE:-4}}
+TOKENIZER=${TOKENIZER:-"Qwen/Qwen3-0.6B"}
+
+INPUT_LEN=${INPUT_LEN:-512}
+OUTPUT_LEN=${OUTPUT_LEN:-128}
+NUM_PROMPTS=${NUM_PROMPTS:-32}
+LOAD_FORMAT=${LOAD_FORMAT:-"dummy"}
+GPU_MEM_UTIL=${GPU_MEM_UTIL:-0.6}
+PYTHON_EXEC=${PYTHON_EXEC:-"python3"}
+
+echo "================================================================="
+echo "Running vLLM Benchmark for MaxText"
+echo "Model Name       : $MODEL_NAME"
+echo "Tokenizer Path   : $TOKENIZER"
+echo "Attention Kernel : $ATTENTION"
+echo "Tensor Parallel  : $TP_SIZE"
+echo "Input / Output   : $INPUT_LEN in / $OUTPUT_LEN out ($NUM_PROMPTS prompts)"
+echo "Load Format      : $LOAD_FORMAT"
+echo "================================================================="
+
+# Set standard TPU & vLLM environment variables
+export PYTHONPATH=$(pwd)/src:${PYTHONPATH}
+export SKIP_JAX_PRECOMPILE=1
+export NEW_MODEL_DESIGN=1
+export VLLM_ENABLE_V1_MULTIPROCESSING=0
+
+if [ "$ATTENTION" = "vllm_batched_rpa" ]; then
+  export USE_BATCHED_RPA_KERNEL=1
+else
+  export USE_BATCHED_RPA_KERNEL=0
+fi
+
+# Run offline throughput benchmark
+$PYTHON_EXEC -m vllm.entrypoints.cli.main bench throughput \
+  --model "$TOKENIZER" \
+  --tensor-parallel-size "$TP_SIZE" \
+  --gpu-memory-utilization "$GPU_MEM_UTIL" \
+  --load-format "$LOAD_FORMAT" \
+  --hf-overrides '{"architectures": ["MaxTextForCausalLM"]}' \
+  --additional-config "{\"maxtext_config\": {\"model_name\": \"$MODEL_NAME\", \"weight_dtype\": \"bfloat16\", \"attention\": \"$ATTENTION\", \"allow_split_physical_axes\": true, \"scan_layers\": true, \"enable_nnx\": true, \"pure_nnx_decoder\": true}}" \
+  --dataset-name random \
+  --random-input-len "$INPUT_LEN" \
+  --random-output-len "$OUTPUT_LEN" \
+  --num-prompts "$NUM_PROMPTS"
+
+echo "================================================================="
+echo "Benchmark completed successfully for $MODEL_NAME ($ATTENTION)!"
+echo "================================================================="