ROCm
diff --git a/‎aiter/ops/triton/attention/mla_decode.py‎
Lines changed: 2 additions & 0 deletions b/‎aiter/ops/triton/attention/mla_decode.py‎
Lines changed: 2 additions & 0 deletions
@@ -53,6 +53,8 @@ def tanh(x):
         "logit_cap",
         "Lk",
         "Lv",
+        "num_warps",
+        "num_stages",
     ],
 )
Original file line number	Diff line number	Diff line change
`@@ -53,6 +53,8 @@ def tanh(x):`
`53`	`53`	`"logit_cap",`
`54`	`54`	`"Lk",`
`55`	`55`	`"Lv",`
	`56`	`+ "num_warps",`
	`57`	`+ "num_stages",`
`56`	`58`	`],`
`57`	`59`	`)`
`58`	`60`