Add experimental attention quantization flags.

This will add two experimental flags to the splash attention config, to
quantize Q and K respectively. Attention quantization is currently a research
project so these should be turned off by default. In the future we will support
more quantization options for the backwards pass, RoPE vs. no-RoPE, more
dtypes, etc.

PiperOrigin-RevId: 949075090

Author: Google-ML-Automation

src/maxtext/configs/base.yml

@@ -1120,6 +1120,8 @@ local_sa_v_layout: None  # inherits from sa_v_layout if None
 local_use_splash_scheduler: None  # inherits from use_splash_scheduler if None
 local_sa_fuse_reciprocal: None  # inherits from sa_fuse_reciprocal if None
 local_sa_use_base2_exp: None  # inherits from sa_use_base2_exp if None
+experimental_sa_quant_q_fp8: False # Experimental flag: If enabled, the Q tensor in splash attention is quantized to jnp.float8_e4m3fn, without scaling factors.
+experimental_sa_quant_k_fp8: False # Experimental flag: If enabled, the K tensor in splash attention is quantized to jnp.float8_e4m3fn, without scaling factors.
 use_max_logit_estimate: -1 # -1 means no estimate, any > 0 value will be used as max logit estimate
 cost_estimate_flops_fwd: -1 # -1 means using splash default cost estmiation, any >= 0 value will be used as cost estmiation for splash to overlap for communication (forward)
 cost_estimate_flops_bwd: -1 # -1 means using splash default cost estmiation, any >= 0 value will be used as cost estmiation for splash to overlap for communication (backward)

src/maxtext/configs/types.py

@@ -740,6 +740,20 @@ class SplashAttention(BaseModel):
   local_use_splash_scheduler: bool | None = Field(None, description="Use experimental local splash attention scheduler.")
   local_sa_fuse_reciprocal: bool | None = Field(None, description="Maps to local fuse_reciprocal in SplashConfig.")
   local_sa_use_base2_exp: bool | None = Field(None, description="Maps to local use_base2_exp in SplashConfig.")
+  experimental_sa_quant_q_fp8: bool | None = Field(
+      None,
+      description=(
+          "Experimental flag: If enabled, the Q tensor in splash attention is"
+          " quantized to jnp.float8_e4m3fn, without scaling factors."
+      ),
+  )
+  experimental_sa_quant_k_fp8: bool | None = Field(
+      None,
+      description=(
+          "Experimental flag: If enabled, the K tensor in splash attention is"
+          " quantized to jnp.float8_e4m3fn, without scaling factors."
+      ),
+  )
   use_max_logit_estimate: int = Field(
       -1,
       description="-1 means no estimate, any > 0 value will be used as max logit estimate",
@@ -3456,6 +3470,16 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
           "Please disable attn_logits_soft_cap when using use_qk_clip."
       )
 
+    if self.experimental_sa_quant_q_fp8 and self.attention_type != "mla":
+      raise ValueError(
+          "Q quantization is currently only supported with"
+          f" attention_type='mla'. Found attention_type='{self.attention_type}'"
+      )
+    if self.experimental_sa_quant_k_fp8 and self.attention_type != "mla":
+      raise ValueError(
+          "K quantization is currently only supported with"
+          f" attention_type='mla'. Found attention_type='{self.attention_type}'"
+      )
     if self.share_kv_projections and self.fused_qkv:
       raise ValueError("`share_kv_projections` is not compatible with `fused_qkv`.")
     if self.share_kv_projections and self.attention_type == "mla":

src/maxtext/layers/attention_mla.py

@@ -876,6 +876,10 @@ def mla_query_projection(
     # Query projection is scaled by self.softmax_scale to be consistent MaxText implementation.
     # DeepSeek v3 was doing it in attention score computation.
     query = jnp.concatenate([q_nope, q_pe], axis=-1) * self.softmax_scale
+
+    if self.config.experimental_sa_quant_q_fp8:
+      query = query.astype(jnp.float8_e4m3fn)
+
     query = self._maybe_shard_with_logical(query, query_logical_name)
     return query, low_rank_q
 
@@ -899,6 +903,9 @@ def mla_get_key_value(self, low_rank_main, key_rope, model_mode):
 
     key = jnp.concatenate([key_nope, key_rope], axis=-1)
 
+    if self.config.experimental_sa_quant_k_fp8:
+      key = key.astype(jnp.float8_e4m3fn)
+
     key = self._maybe_shard_with_logical(key, key_logical_name)
     value = self._maybe_shard_with_logical(value, value_logical_name)
     return key, value