Guard Stride tl.assume

Summary: Guards the tl.assume values generated by the strides behind a non-negative requirement. This is not due to any regression, just a reviewer suggestion.

Differential Revision: D72398678

Author: njriasan

Diff for: fbgemm_gpu/experimental/gemm/triton_gemm/fp8_gemm.py

@@ -1215,6 +1215,12 @@ def persistent_grid(META):
             # GROUP_M=8,
             # USE_BIAS=bias is not None,
             AB_DTYPE=False,
+            NON_NEGATIVE_STRIDE_AM=a.stride(0) >= 0,
+            NON_NEGATIVE_STRIDE_AK=a.stride(1) >= 0,
+            NON_NEGATIVE_STRIDE_BN=b.stride(0) >= 0,
+            NON_NEGATIVE_STRIDE_BK=b.stride(1) >= 0,
+            NON_NEGATIVE_STRIDE_CM=c.stride(0) >= 0,
+            NON_NEGATIVE_STRIDE_CN=c.stride(1) >= 0,
         )
     elif use_warp_specialization:
         assert has_warp_specialization
@@ -3363,6 +3369,12 @@ def _kernel_matmul_fp8_row_non_persistent(
     SPLIT_K: tl.constexpr,
     EVEN_K: tl.constexpr,
     AB_DTYPE: tl.constexpr,
+    NON_NEGATIVE_STRIDE_AM: tl.constexpr,
+    NON_NEGATIVE_STRIDE_AK: tl.constexpr,
+    NON_NEGATIVE_STRIDE_BN: tl.constexpr,
+    NON_NEGATIVE_STRIDE_BK: tl.constexpr,
+    NON_NEGATIVE_STRIDE_CM: tl.constexpr,
+    NON_NEGATIVE_STRIDE_CN: tl.constexpr,
 ) -> None:
     """Matmul kernel of [M, K] @ [N, K] with row-wise scales
 
@@ -3400,12 +3412,18 @@ def _kernel_matmul_fp8_row_non_persistent(
     tl.assume(M >= 0)
     tl.assume(N >= 0)
     tl.assume(K >= 0)
-    tl.assume(stride_am >= 0)
-    tl.assume(stride_ak >= 0)
-    tl.assume(stride_bn >= 0)
-    tl.assume(stride_bk >= 0)
-    tl.assume(stride_cm >= 0)
-    tl.assume(stride_cn >= 0)
+    if NON_NEGATIVE_STRIDE_AM:
+        tl.assume(stride_am >= 0)
+    if NON_NEGATIVE_STRIDE_AK:
+        tl.assume(stride_ak >= 0)
+    if NON_NEGATIVE_STRIDE_BN:
+        tl.assume(stride_bn >= 0)
+    if NON_NEGATIVE_STRIDE_BK:
+        tl.assume(stride_bk >= 0)
+    if NON_NEGATIVE_STRIDE_CM:
+        tl.assume(stride_cm >= 0)
+    if NON_NEGATIVE_STRIDE_CN:
+        tl.assume(stride_cn >= 0)
     # Matrix multiplication.
     pid = tl.program_id(0)
     pid_z = tl.program_id(1)