Enable experimental V quantization, quantization strategies.

src/maxtext/configs/base.yml

@@ -1125,8 +1125,9 @@ 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.
+experimental_sa_quant_q_fp8: "none" # Experimental flag: Quantization strategy for Q in splash attention ("none", "cast", or "tensor").
+experimental_sa_quant_k_fp8: "none" # Experimental flag: Quantization strategy for K in splash attention ("none", "cast", or "tensor").
+experimental_sa_quant_v_fp8: "none" # Experimental flag: Quantization strategy for V in splash attention ("none", "cast", or "tensor").
 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

@@ -720,6 +720,36 @@ class Llama4Attention(BaseModel):
       description="Dynamically scale attention temperature based on sequence length.",
   )
 
+class ExperimentalAttentionQuantizationType(str, Enum):
+  """Experimental scaling strategies for attention quantization."""
+  NONE = "none"
+  CAST = "cast" # Tensor is cast to jnp.float8_e4m3fn without scaling factors.
+  TENSOR = "tensor"
+
+
+class ExperimentalAttentionQuantization(BaseModel):
+  """Experimental options for attention quantization."""
+  experimental_sa_quant_q_fp8: ExperimentalAttentionQuantizationType | None = Field(
+      ExperimentalAttentionQuantizationType.NONE,
+      description=(
+          "Experimental flag: If enabled, the Q tensor in splash attention is"
+          " quantized to jnp.float8_e4m3fn."
+      ),
+  )
+  experimental_sa_quant_k_fp8: ExperimentalAttentionQuantizationType | None = Field(
+      ExperimentalAttentionQuantizationType.NONE,
+      description=(
+          "Experimental flag: If enabled, the K tensor in splash attention is"
+          " quantized to jnp.float8_e4m3fn."
+      ),
+  )
+  experimental_sa_quant_v_fp8: ExperimentalAttentionQuantizationType | None = Field(
+      ExperimentalAttentionQuantizationType.NONE,
+      description=(
+          "Experimental flag: If enabled, the V tensor in splash attention is"
+          " quantized to jnp.float8_e4m3fn."
+      ),
+  )
 
 class SplashAttention(BaseModel):
   """Tunable block sizes for Splash Attention kernels."""
@@ -759,20 +789,6 @@ 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",
@@ -2543,6 +2559,7 @@ class MaxTextConfig(
     MTP,
     LogitsAndLoss,
     # Attention Mechanisms
+    ExperimentalAttentionQuantization,
     Attention,
     MlaAttention,
     CompressedAttention,
@@ -3569,16 +3586,33 @@ 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":
+    if (
+        self.experimental_sa_quant_q_fp8 is not None
+        and self.experimental_sa_quant_q_fp8 != ExperimentalAttentionQuantizationType.NONE
+        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":
+    if (
+        self.experimental_sa_quant_k_fp8 is not None
+        and self.experimental_sa_quant_k_fp8 != ExperimentalAttentionQuantizationType.NONE
+        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.experimental_sa_quant_v_fp8 is not None
+        and self.experimental_sa_quant_v_fp8 != ExperimentalAttentionQuantizationType.NONE
+        and self.attention_type != "mla"
+    ):
+      raise ValueError(
+          "V 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

@@ -896,8 +896,10 @@ def mla_query_projection(
     # 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:
+    if self.config.experimental_sa_quant_q_fp8 == "cast":
       query = query.astype(jnp.float8_e4m3fn)
+    elif self.config.experimental_sa_quant_q_fp8 == "tensor":
+      raise NotImplementedError("Tensor scaling for Q is not implemented yet.")
 
     query = self._maybe_shard_with_logical(query, query_logical_name)
     return query, low_rank_q
@@ -922,8 +924,14 @@ 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:
+    if self.config.experimental_sa_quant_k_fp8 == "cast":
       key = key.astype(jnp.float8_e4m3fn)
+    elif self.config.experimental_sa_quant_k_fp8 == "tensor":
+      raise NotImplementedError("Tensor scaling for K is not implemented yet.")
+    if self.config.experimental_sa_quant_v_fp8 == "cast":
+      value = value.astype(jnp.float8_e4m3fn)
+    elif self.config.experimental_sa_quant_v_fp8 == "tensor":
+      raise NotImplementedError("Tensor scaling for V is not implemented yet.")
 
     key = self._maybe_shard_with_logical(key, key_logical_name)
     value = self._maybe_shard_with_logical(value, value_logical_name)