Multi-Head Latent Attention (MLA) layer.
| 210 | |
| 211 | |
| 212 | class MLA(Attention): |
| 213 | """Multi-Head Latent Attention (MLA) layer.""" |
| 214 | |
| 215 | def __init__( |
| 216 | self, |
| 217 | config: Config, |
| 218 | num_query_heads: int, |
| 219 | num_kv_heads: int, |
| 220 | head_dim: int, |
| 221 | max_target_length: int, |
| 222 | mesh: Mesh, |
| 223 | attention_kernel: str, |
| 224 | inputs_q_shape: Tuple, |
| 225 | inputs_kv_shape: Tuple, |
| 226 | dtype: DType = jnp.float32, |
| 227 | weight_dtype: DType = jnp.float32, |
| 228 | max_prefill_predict_length: int = -1, |
| 229 | dropout_rate: float = 0.0, |
| 230 | kernel_init: NdInitializer = nd_dense_init(1.0, "fan_in", "normal"), |
| 231 | float32_qk_product: bool = False, # computes logits in float32 for stability. |
| 232 | float32_logits: bool = False, # cast logits in float32 for stability. |
| 233 | quant: Optional[Quant] = None, |
| 234 | kv_quant: Optional[KVQuant] = None, |
| 235 | attention_type: AttentionType = AttentionType.MLA, # Default to MLA attention |
| 236 | attn_logits_soft_cap: float | None = None, |
| 237 | sliding_window_size: int | None = None, |
| 238 | use_ragged_attention: bool = False, |
| 239 | ragged_block_size: int = 256, |
| 240 | use_qk_norm: bool = False, |
| 241 | query_pre_attn_scalar: float | None = None, |
| 242 | use_bias_in_projections: bool = False, # Set to True will enable bias in q, k, v, o projections |
| 243 | # Temperature tuning parameters used for Llama4 |
| 244 | temperature_tuning: bool = False, |
| 245 | temperature_tuning_scale: float = 0.1, |
| 246 | temperature_tuning_floor_scale: float = 8192.0, |
| 247 | # Shard the query activation as the same as the key and value. |
| 248 | # TODO: Find a better sharding axis name. |
| 249 | # TODO: Further break down the Training and Inference axes for the q, k, v. |
| 250 | prefill_query_axis_names: AxisNames = (PREFILL_KV_BATCH, PREFILL_LENGTH, KV_HEAD, KV_HEAD_DIM), |
| 251 | prefill_key_axis_names: AxisNames = (PREFILL_KV_BATCH, PREFILL_LENGTH, KV_HEAD, KV_HEAD_DIM), |
| 252 | prefill_value_axis_names: AxisNames = (PREFILL_KV_BATCH, PREFILL_LENGTH, KV_HEAD, KV_HEAD_DIM), |
| 253 | query_axis_names: AxisNames = (KV_BATCH, LENGTH_NO_EXP, KV_HEAD, KV_HEAD_DIM), |
| 254 | key_axis_names: AxisNames = (KV_BATCH, LENGTH_NO_EXP, KV_HEAD, KV_HEAD_DIM), |
| 255 | value_axis_names: AxisNames = (KV_BATCH, LENGTH_NO_EXP, KV_HEAD, KV_HEAD_DIM), |
| 256 | ep_query_axis_names: AxisNames = (KV_BATCH_NO_EXP, LENGTH, KV_HEAD, KV_HEAD_DIM), |
| 257 | ep_key_axis_names: AxisNames = (KV_BATCH_NO_EXP, LENGTH, KV_HEAD, KV_HEAD_DIM), |
| 258 | ep_value_axis_names: AxisNames = (KV_BATCH_NO_EXP, LENGTH, KV_HEAD, KV_HEAD_DIM), |
| 259 | input_axis_names: AxisNames = (BATCH, LENGTH_NO_EXP, EMBED), |
| 260 | ep_input_axis_names: AxisNames = (BATCH_NO_EXP, LENGTH, EMBED), |
| 261 | out_axis_names: AxisNames = (BATCH, LENGTH_NO_EXP, HEAD, D_KV), |
| 262 | ep_out_axis_names: AxisNames = (BATCH_NO_EXP, LENGTH, HEAD, D_KV), |
| 263 | prefill_input_axis_names: AxisNames = (PREFILL_KV_BATCH, PREFILL_LENGTH, EMBED), |
| 264 | decode_input_axis_names: AxisNames = (DECODE_BATCH, DECODE_LENGTH, EMBED), |
| 265 | prefill_out_axis_names: AxisNames = (PREFILL_KV_BATCH, PREFILL_LENGTH, HEAD, D_KV), |
| 266 | decode_out_axis_names: AxisNames = (DECODE_BATCH, DECODE_LENGTH, HEAD, D_KV), |
| 267 | prefill_cache_axis_order: AxisIdxes = (1, 2, 0, 3), |
| 268 | ar_cache_axis_order: AxisIdxes = (1, 2, 0, 3), |
| 269 | compute_axis_order: AxisIdxes = (0, 1, 2, 3), |
no outgoing calls