Merge pull request #3046 from AI-Hypercomputer:zxhe/add-attn-soft-cap

Google-ML-Automation · Google-ML-Automation · commit 04c8b12c9d77 · 2026-01-29T14:08:36.000-08:00
PiperOrigin-RevId: 862893155
diff --git a/src/MaxText/layers/llama2.py b/src/MaxText/layers/llama2.py
@@ -101,6 +101,7 @@ def __init__(
         use_ragged_attention=config.use_ragged_attention,
         ragged_block_size=config.ragged_block_size,
         model_mode=model_mode,
+        attn_logits_soft_cap=config.attn_logits_soft_cap,
         rngs=rngs,
     )
 

Original file line number	Diff line number	Diff line change
`@@ -101,6 +101,7 @@ def __init__(`
`101`	`101`	`use_ragged_attention=config.use_ragged_attention,`
`102`	`102`	`ragged_block_size=config.ragged_block_size,`
`103`	`103`	`model_mode=model_mode,`
	`104`	`+ attn_logits_soft_cap=config.attn_logits_soft_cap,`
`104`	`105`	`rngs=rngs,`
`105`	`106`	`)`
`106`	`107`