remove decode mode in rpa

Author: JamesBrianD

benchmark/kernels/flash_attention/bench_flashattention.py

@@ -112,7 +112,6 @@ def jitted_attn(
             cu_kv_lens,
             distribution,
             custom_mask=None,
-            decode_mode=0,
             causal=1,
             sm_scale=sm_scale,
         )

benchmark/kernels/flash_attention/get_block_spec_config.py

@@ -105,7 +105,6 @@ def jitted_attn(
             cu_kv_lens,
             distribution,
             None,
-            decode_mode=0,
             sm_scale=sm_scale,
             num_kv_pages_per_block=num_kv_pages_per_block,
             num_queries_per_block=num_queries_per_block,

eplb/256expert_16chip_fp8_random_20260210_175806.npy

@@ -18,7 +18,6 @@
 from jax.experimental.pallas import tpu as pltpu
 
 from sgl_jax.srt.kernels.ragged_paged_attention.tuned_block_sizes import (
-    get_kv_pages_for_decode,
     get_tuned_block_sizes,
 )
 from sgl_jax.srt.kernels.ragged_paged_attention.util import (
@@ -1381,7 +1380,6 @@ def get_kernel_scope_name(bq_size, bkv_p, page_size):
     jax.jit,
     static_argnames=(
         "causal",
-        "decode_mode",
         "sm_scale",
         "sliding_window",
         "soft_cap",
@@ -1410,7 +1408,6 @@ def ragged_paged_attention(
     custom_mask: jax.Array,  # if causal is True, custom_mask shape is [patten_total_kv_len], else [0]
     *,
     causal: int = 1,  # 1: True, 0: False
-    decode_mode: int = 1,
     sm_scale: float = 1.0,
     sliding_window: int | None = None,
     soft_cap: float | None = None,
@@ -1511,29 +1508,17 @@ def ragged_paged_attention(
     bkv_p = num_kv_pages_per_block
     bq_sz = num_queries_per_block
     if bq_sz is None or bkv_p is None:
-        if decode_mode == 1:
-            bkv_p = get_kv_pages_for_decode(
-                q.dtype,
-                kv_cache_fused_processed.dtype,
-                actual_num_q_heads,
-                actual_num_kv_heads,
-                head_dim,
-                page_size,
-                pages_per_seq,
-            )
-            bq_sz = 1
-        else:
-            bkv_p, bq_sz = get_tuned_block_sizes(
-                q.dtype,
-                kv_cache_fused_processed.dtype,
-                actual_num_q_heads,
-                actual_num_kv_heads,
-                head_dim,
-                page_size,
-                max_num_tokens,
-                pages_per_seq,
-                causal,
-            )
+        bkv_p, bq_sz = get_tuned_block_sizes(
+            q.dtype,
+            kv_cache_fused_processed.dtype,
+            actual_num_q_heads,
+            actual_num_kv_heads,
+            head_dim,
+            page_size,
+            max_num_tokens,
+            pages_per_seq,
+            causal,
+        )
     kv_packing = get_dtype_packing(kv_cache_fused_processed.dtype)
     if page_size == 1:
         bkv_p = bkv_p // 2

eplb/expert_dist_orz_math_10000_steps.npy

@@ -1514,195 +1514,3 @@ def get_simplified_key(
         next_power_of_2(page_size),
         next_power_of_2(max_num_tokens),
     )
-
-
-TUNED_KV_PAGES_FOR_DECODE = {
-    # key
-    #   - device_name
-    #     - q dtype
-    #     - kv dtype
-    #     - q head number
-    #     - kv head number
-    #     - head dim
-    #     - page_size
-    # value:
-    #   - num_kv_pages_per_block
-    "TPU v6e": {
-        ("bfloat16", "bfloat16", 1, 1, 128, 128): 256,
-        ("bfloat16", "bfloat16", 2, 1, 128, 128): 256,
-        ("bfloat16", "bfloat16", 4, 1, 128, 128): 256,
-        ("bfloat16", "bfloat16", 8, 1, 128, 128): 256,
-        ("bfloat16", "bfloat16", 16, 1, 128, 128): 256,
-        ("bfloat16", "bfloat16", 2, 2, 128, 128): 128,
-        ("bfloat16", "bfloat16", 4, 2, 128, 128): 128,
-        ("bfloat16", "bfloat16", 8, 2, 128, 128): 128,
-        ("bfloat16", "bfloat16", 16, 2, 128, 128): 128,
-        ("bfloat16", "bfloat16", 32, 2, 128, 128): 128,
-        ("bfloat16", "bfloat16", 4, 4, 128, 128): 64,
-        ("bfloat16", "bfloat16", 8, 4, 128, 128): 64,
-        ("bfloat16", "bfloat16", 16, 4, 128, 128): 64,
-        ("bfloat16", "bfloat16", 32, 4, 128, 128): 64,
-        ("bfloat16", "bfloat16", 8, 8, 128, 128): 32,
-        ("bfloat16", "bfloat16", 16, 8, 128, 128): 32,
-        ("bfloat16", "bfloat16", 32, 8, 128, 128): 32,
-        ("bfloat16", "bfloat16", 64, 8, 128, 128): 32,
-        ("bfloat16", "bfloat16", 16, 16, 128, 128): 16,
-        ("bfloat16", "bfloat16", 32, 16, 128, 128): 16,
-        ("bfloat16", "bfloat16", 64, 16, 128, 128): 16,
-        ("bfloat16", "bfloat16", 128, 16, 128, 128): 16,
-        ("bfloat16", "bfloat16", 1, 1, 128, 256): 128,
-        ("bfloat16", "bfloat16", 2, 1, 128, 256): 128,
-        ("bfloat16", "bfloat16", 4, 1, 128, 256): 128,
-        ("bfloat16", "bfloat16", 8, 1, 128, 256): 128,
-        ("bfloat16", "bfloat16", 16, 1, 128, 256): 128,
-        ("bfloat16", "bfloat16", 2, 2, 128, 256): 64,
-        ("bfloat16", "bfloat16", 4, 2, 128, 256): 64,
-        ("bfloat16", "bfloat16", 8, 2, 128, 256): 64,
-        ("bfloat16", "bfloat16", 16, 2, 128, 256): 64,
-        ("bfloat16", "bfloat16", 4, 4, 128, 256): 32,
-        ("bfloat16", "bfloat16", 8, 4, 128, 256): 32,
-        ("bfloat16", "bfloat16", 16, 4, 128, 256): 32,
-        ("bfloat16", "bfloat16", 32, 4, 128, 256): 32,
-        ("bfloat16", "bfloat16", 8, 8, 128, 256): 16,
-        ("bfloat16", "bfloat16", 16, 8, 128, 256): 16,
-        ("bfloat16", "bfloat16", 32, 8, 128, 256): 16,
-        ("bfloat16", "bfloat16", 64, 8, 128, 256): 16,
-        ("bfloat16", "bfloat16", 16, 16, 128, 256): 8,
-        ("bfloat16", "bfloat16", 32, 16, 128, 256): 8,
-        ("bfloat16", "bfloat16", 64, 16, 128, 256): 8,
-        ("bfloat16", "bfloat16", 128, 16, 128, 256): 8,
-        ("bfloat16", "bfloat16", 256, 16, 128, 256): 8,
-        ("bfloat16", "bfloat16", 512, 16, 128, 256): 8,
-    },
-    "TPU v7": {
-        ("bfloat16", "bfloat16", 1, 1, 128, 128): 256,
-        ("bfloat16", "bfloat16", 2, 1, 128, 128): 256,
-        ("bfloat16", "bfloat16", 4, 1, 128, 128): 256,
-        ("bfloat16", "bfloat16", 8, 1, 128, 128): 256,
-        ("bfloat16", "bfloat16", 16, 1, 128, 128): 256,
-        ("bfloat16", "bfloat16", 2, 2, 128, 128): 128,
-        ("bfloat16", "bfloat16", 4, 2, 128, 128): 128,
-        ("bfloat16", "bfloat16", 8, 2, 128, 128): 128,
-        ("bfloat16", "bfloat16", 16, 2, 128, 128): 128,
-        ("bfloat16", "bfloat16", 32, 2, 128, 128): 128,
-        ("bfloat16", "bfloat16", 4, 4, 128, 128): 64,
-        ("bfloat16", "bfloat16", 8, 4, 128, 128): 64,
-        ("bfloat16", "bfloat16", 16, 4, 128, 128): 64,
-        ("bfloat16", "bfloat16", 32, 4, 128, 128): 64,
-        ("bfloat16", "bfloat16", 8, 8, 128, 128): 32,
-        ("bfloat16", "bfloat16", 16, 8, 128, 128): 32,
-        ("bfloat16", "bfloat16", 32, 8, 128, 128): 32,
-        ("bfloat16", "bfloat16", 64, 8, 128, 128): 32,
-        ("bfloat16", "bfloat16", 16, 16, 128, 128): 16,
-        ("bfloat16", "bfloat16", 32, 16, 128, 128): 16,
-        ("bfloat16", "bfloat16", 64, 16, 128, 128): 16,
-        ("bfloat16", "bfloat16", 128, 16, 128, 128): 16,
-        ("bfloat16", "bfloat16", 1, 1, 128, 256): 128,
-        ("bfloat16", "bfloat16", 2, 1, 128, 256): 128,
-        ("bfloat16", "bfloat16", 4, 1, 128, 256): 128,
-        ("bfloat16", "bfloat16", 8, 1, 128, 256): 128,
-        ("bfloat16", "bfloat16", 16, 1, 128, 256): 128,
-        ("bfloat16", "bfloat16", 2, 2, 128, 256): 64,
-        ("bfloat16", "bfloat16", 4, 2, 128, 256): 64,
-        ("bfloat16", "bfloat16", 8, 2, 128, 256): 64,
-        ("bfloat16", "bfloat16", 16, 2, 128, 256): 64,
-        ("bfloat16", "bfloat16", 4, 4, 128, 256): 32,
-        ("bfloat16", "bfloat16", 8, 4, 128, 256): 32,
-        ("bfloat16", "bfloat16", 16, 4, 128, 256): 32,
-        ("bfloat16", "bfloat16", 32, 4, 128, 256): 32,
-        ("bfloat16", "bfloat16", 8, 8, 128, 256): 16,
-        ("bfloat16", "bfloat16", 16, 8, 128, 256): 16,
-        ("bfloat16", "bfloat16", 32, 8, 128, 256): 16,
-        ("bfloat16", "bfloat16", 64, 8, 128, 256): 16,
-        ("bfloat16", "bfloat16", 16, 16, 128, 256): 8,
-        ("bfloat16", "bfloat16", 32, 16, 128, 256): 8,
-        ("bfloat16", "bfloat16", 64, 16, 128, 256): 8,
-        ("bfloat16", "bfloat16", 128, 16, 128, 256): 8,
-        ("bfloat16", "bfloat16", 256, 16, 128, 256): 8,
-        ("bfloat16", "bfloat16", 512, 16, 128, 256): 8,
-    },
-}
-
-
-def get_kv_pages_for_decode(
-    q_dtype,
-    kv_dtype,
-    actual_num_q_heads,
-    actual_num_kv_heads,
-    head_dim,
-    page_size,
-    pages_per_seq,
-    causal=True,
-) -> int:
-    if not causal:
-        # FIXME(pc) hack this to avoid oom when precompile, currently, we still have no better choice for non-causal's mask
-        # this should be optimied future
-        return 4
-    """Look up for the best num_kv_pages_per_blk from auto-tuned table."""
-    tpu_version = get_tpu_version()
-
-    if tpu_version < 4:
-        raise NotImplementedError("TPU version must be 4 or higher.")
-    keys = get_simplified_key_for_decode(
-        page_size,
-        q_dtype,
-        kv_dtype,
-        actual_num_q_heads,
-        actual_num_kv_heads,
-        head_dim,
-    )
-
-    device_name = keys[0]
-
-    # Default block sizes.
-    bkv_p = 1024 // page_size
-    if tpu_version == 4:
-        # TPUv4 has much smaller VMEM size so we pick fixed block sizes.
-        bkv_p = 512 // page_size
-    else:
-        if (
-            device_name in TUNED_KV_PAGES_FOR_DECODE
-            and keys[1:] in TUNED_KV_PAGES_FOR_DECODE[device_name]
-        ):
-            bkv_p = TUNED_KV_PAGES_FOR_DECODE[device_name][keys[1:]]
-        else:
-            logger.info(
-                "Tuned RPA kv page not found for %s: page_size=%s, actual_num_q_heads=%s, "
-                "actual_num_kv_heads=%s, head_dim=%s, pages_per_seq=%s.",
-                device_name,
-                page_size,
-                actual_num_q_heads,
-                actual_num_kv_heads,
-                head_dim,
-                pages_per_seq,
-            )
-            logger.info("Using default block size: bkv_p=%s.", bkv_p)
-
-    return min(pages_per_seq, bkv_p)
-
-
-def get_simplified_key_for_decode(
-    page_size,
-    q_dtype,
-    kv_dtype,
-    num_q_heads,
-    num_kv_heads,
-    head_dim,
-):
-    """Get the simplified key to reduce the number of combinations."""
-    assert num_q_heads % num_kv_heads == 0
-    device = get_device_name()
-    q_dtype_name = jnp.dtype(q_dtype).name
-    kv_dtype_name = jnp.dtype(kv_dtype).name
-    num_q_heads = next_power_of_2(num_q_heads)
-    num_kv_heads = next_power_of_2(num_kv_heads)
-
-    return (
-        device,
-        q_dtype_name,
-        kv_dtype_name,
-        num_q_heads,
-        num_kv_heads,
-        (head_dim + 127) // 128 * 128,
-        next_power_of_2(page_size),
-    )

eplb/expert_dist_orz_math_5000_steps.npy

@@ -510,8 +510,6 @@ def __call__(
         if hasattr(token_to_kv_pool, "remap_cache_loc") and self.page_size == 1:
             page_indices_arg = token_to_kv_pool.remap_cache_loc(page_indices_arg, layer.layer_id)
 
-        decode_mode = 1 if forward_batch.forward_mode == ForwardMode.DECODE else 0
-
         in_specs = (
             P(self.attention_data_partition_axis, self.kv_partition_axis),  # queries
             P(self.attention_data_partition_axis, self.kv_partition_axis),  # keys (new tokens)
@@ -545,7 +543,6 @@ def _ragged_paged_attention_with_fused_kv(*args):
                 kv_cache_fused,
                 *other_args,
                 causal=causal,
-                decode_mode=decode_mode,
                 sm_scale=scale,
                 sliding_window=layer.sliding_window_size,
                 soft_cap=layer.logit_cap,