[KERNELS] fix handling of opt_flags.idle_sms (#7211)

ptillet · web-flow · commit 0e9706cd4f5d · 2025-06-17T20:44:19.000-07:00
diff --git a/python/triton_kernels/triton_kernels/matmul_ogs.py b/python/triton_kernels/triton_kernels/matmul_ogs.py
@@ -705,7 +705,7 @@ def matmul_ogs(x, w, bias,
     )
     # matrix multiplication
     n_cta = batch_size * grid_m * grid_n * opt_flags.split_k
-    n_cta = min(target_info.num_sms(), n_cta) if opt_flags.is_persistent else n_cta
+    n_cta = min(target_info.num_sms() - opt_flags.idle_sms, n_cta) if opt_flags.is_persistent else n_cta
     flex = precision_config.flex_ctx
     bias_stride = None if bias is None else bias.stride(0)
     num_indx = None if scatter_indx is None else scatter_indx.src_indx.shape[0]

Original file line number	Diff line number	Diff line change
`@@ -705,7 +705,7 @@ def matmul_ogs(x, w, bias,`
`705`	`705`	`)`
`706`	`706`	`# matrix multiplication`
`707`	`707`	`n_cta = batch_size * grid_m * grid_n * opt_flags.split_k`
`708`		`- n_cta = min(target_info.num_sms(), n_cta) if opt_flags.is_persistent else n_cta`
	`708`	`+ n_cta = min(target_info.num_sms() - opt_flags.idle_sms, n_cta) if opt_flags.is_persistent else n_cta`
`709`	`709`	`flex = precision_config.flex_ctx`
`710`	`710`	`bias_stride = None if bias is None else bias.stride(0)`
`711`	`711`	`num_indx = None if scatter_indx is None else scatter_indx.src_indx.shape[0]`