Address comments

Signed-off-by: Harsh Menon <[email protected]>
iree-org · Feb 10, 2025 · b20eb49 · b20eb49
1 parent 0bd1e23
commit b20eb49
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/iree/turbine/kernel/wave/templates/extend_attention.py b/iree/turbine/kernel/wave/templates/extend_attention.py
@@ -71,7 +71,7 @@ def get_extend_attention_kernel(
     LOG2E = 1.44269504089
     logit_cap *= LOG2E
     dk_sqrt = math.sqrt(1.0 / shape.head_size)
-    layer_scaling = layer_scaling or dk_sqrt * LOG2E
+    layer_scaling = (layer_scaling or dk_sqrt) * LOG2E
 
     constraints: list[tkw.Constraint] = []
     constraints += [