Revert the parameter selection since it is adversarial on M4.

liuliu · Dec 17, 2024 · a16537f · a16537f
1 parent 3bff50d
commit a16537f
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/lib/nnc/mfa/v2/AttentionDescriptor.cpp b/lib/nnc/mfa/v2/AttentionDescriptor.cpp
@@ -454,10 +454,10 @@ std::vector<AttentionParameterRow> AttentionDescriptor::forwardMixed(MTL::Device
   if (device->supportsFamily(MTL::GPUFamily(1009))) {
     return {
       AttentionParameterRow(32, 16, 128, 16, { AttentionOperand::Q, AttentionOperand::O }),
-      AttentionParameterRow(64, 16, 128, 32, { AttentionOperand::Q, AttentionOperand::O }),
-      AttentionParameterRow(160, 32, 128, 32, { AttentionOperand::O }),
-      AttentionParameterRow(224, 32, 128, 32, { AttentionOperand::Q }),
-      AttentionParameterRow(384, 32, 128, 32, {})
+      AttentionParameterRow(96, 16, 128, 32, { AttentionOperand::Q, AttentionOperand::O }),
+      AttentionParameterRow(160, 16, 128, 32, { AttentionOperand::O }),
+      AttentionParameterRow(224, 16, 128, 32, { AttentionOperand::Q }),
+      AttentionParameterRow(384, 16, 128, 32, {})
     };
   } else {
     return {