[KERNELS] added option and test to set idle sms in matmul_ogs (#7210)

Author: ptillet

python/triton_kernels/tests/test_matmul.py

@@ -9,7 +9,7 @@
 import triton_kernels.matmul_ogs_details.opt_flags as opt_flags
 from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig, MicroscalingCtx, FusedActivation, FnSpecs
 from triton_kernels.matmul_ogs import can_use_persistent_tma
-from triton_kernels.matmul_ogs import matmul_ogs, matmul_ogs_torch
+from triton_kernels.matmul_ogs import matmul_ogs_set_idle_sms, matmul_ogs, matmul_ogs_torch
 from triton_kernels.swiglu import swiglu, swiglu_fn, PrecisionConfig as SwiGLUPrecisionConfig
 # numerics utilities
 from triton_kernels.numerics import InFlexData, OutFlexData
@@ -444,6 +444,16 @@ def round_x(x, idx):
                 tri_y_scale).abs() < 1e-10, f"ref_y_scale: {ref_y_scale}, tri_y_scale: {tri_y_scale.item()}"
 
 
+def test_set_idle_sms():
+    from triton_kernels.matmul_ogs_details.opt_flags import make_opt_flags
+    num_idle_sms = 24
+    matmul_ogs_set_idle_sms(num_idle_sms)
+    flags = make_opt_flags(torch.float32, torch.float32, torch.float32, PrecisionConfig(), 0, 0, 0, None, False, False,
+                           1)
+    assert flags.is_persistent
+    assert flags.idle_sms == num_idle_sms
+
+
 @pytest.mark.parametrize("m, n, k, mode", [
     (1200, 704, 608, "ragged"),
     (800, 800, 400, "batched"),

python/triton_kernels/triton_kernels/matmul_ogs.py

@@ -13,7 +13,7 @@
 from .matmul_ogs_details._matmul_ogs import _matmul_ogs
 from .matmul_ogs_details._p_matmul_ogs import _p_matmul_ogs, get_per_device_per_stream_alloc_fn
 from .matmul_ogs_details._finalize_matmul import _finalize_matmul
-from .matmul_ogs_details.opt_flags import make_opt_flags, OptFlags
+from .matmul_ogs_details.opt_flags import make_opt_flags, OptFlags, update_opt_flags_constraints
 from .matmul_ogs_details.fast_contiguous import fast_contiguous
 from .numerics_details.mxfp import SwizzlingType
 from .specialize import specialize
@@ -613,6 +613,8 @@ def _create_tma_descriptors(
 
     return x_tensor_or_desc, w_desc_and_transpose, mx_desc_and_transpose
 
+def matmul_ogs_set_idle_sms(num_idle_sms):
+    update_opt_flags_constraints({"is_persistent": True, "idle_sms": num_idle_sms})
 
 def matmul_ogs(x, w, bias,
                routing_data: RoutingData | None = None,