Fix ops backward compatibility tests

tests/test_flash_attn.py

@@ -20,7 +20,8 @@ from test_util import (
 import kernels
 
 flash_attn3 = kernels.get_kernel("kernels-community/flash-attn3")
-ops = flash_attn3._ops
+ops = flash_attn3._ops.ops
+add_op_namespace_prefix = flash_attn3._ops.add_op_namespace_prefix
 
 
 DISABLE_BACKWARD = os.getenv("FLASH_ATTENTION_DISABLE_BACKWARD", "FALSE") == "TRUE"
@@ -1135,7 +1136,7 @@ def test_flash3_bw_compatibility() -> None:
     # 1/ Instead of removing arguments, error out if their value is no longer supported
     # 2/ When adding arguments, add them at the end with a default value
     assert ops.fwd.default._schema.is_backward_compatible_with(parse_schema(
-        "flash_attn_3::fwd(Tensor q, Tensor k, Tensor v, Tensor(k_new!)? k_new=None, "
+        add_op_namespace_prefix("fwd(Tensor q, Tensor k, Tensor v, Tensor(k_new!)? k_new=None, "
         "Tensor(v_new!)? v_new=None, Tensor? q_v=None, Tensor(out!)? out=None, "
         "Tensor? cu_seqlens_q=None, Tensor? cu_seqlens_k=None, "
         "Tensor? cu_seqlens_k_new=None, Tensor? seqused_q=None, Tensor? seqused_k=None, "
@@ -1146,25 +1147,25 @@ def test_flash3_bw_compatibility() -> None:
         "int attention_chunk=0, float softcap=0., bool is_rotary_interleaved=False, "
         "Tensor? scheduler_metadata=None, int num_splits=0, bool? pack_gqa=None, int sm_margin=0) "
         "-> (Tensor(out!), Tensor, Tensor, Tensor)"
-    ))
+    )))
     assert ops.bwd.default._schema.is_backward_compatible_with(parse_schema(
-        "flash_attn_3::bwd(Tensor dout, Tensor q, Tensor k, Tensor v, Tensor out, Tensor softmax_lse, "
+        add_op_namespace_prefix("bwd(Tensor dout, Tensor q, Tensor k, Tensor v, Tensor out, Tensor softmax_lse, "
         "Tensor(dq!)? dq=None, Tensor(dk!)? dk=None, Tensor(dv!)? dv=None, Tensor? cu_seqlens_q=None, "
         "Tensor? cu_seqlens_k=None, Tensor? seqused_q=None, Tensor? seqused_k=None, int? max_seqlen_q=None, "
         "int? max_seqlen_k=None, float? softmax_scale=None, bool is_causal=False, int window_size_left=-1, "
         "int window_size_right=-1, float softcap=0., bool deterministic=False, int sm_margin=0) "
         "-> (Tensor(dq!), Tensor(dk!), Tensor(dv!), Tensor, Tensor, Tensor, Tensor, Tensor)"
-    ))
+    )))
     assert ops.fwd_combine.default._schema.is_backward_compatible_with(parse_schema(
-        "flash_attn_3::fwd_combine(Tensor out_partial, Tensor lse_partial, Tensor(out!)? out=None, "
+        add_op_namespace_prefix("fwd_combine(Tensor out_partial, Tensor lse_partial, Tensor(out!)? out=None, "
         "ScalarType? out_dtype=None) -> (Tensor(out!), Tensor)"
-    ))
+    )))
     assert ops.get_scheduler_metadata.default._schema.is_backward_compatible_with(parse_schema(
-        "flash_attn_3::get_scheduler_metadata(int batch_size, int max_seqlen_q, int max_seqlen_k, "
+        add_op_namespace_prefix("get_scheduler_metadata(int batch_size, int max_seqlen_q, int max_seqlen_k, "
         "int num_heads, int num_heads_k, int headdim, int headdim_v, ScalarType qkv_dtype, Tensor seqused_k, "
         "Tensor? cu_seqlens_q=None, Tensor? cu_seqlens_k=None, Tensor? cu_seqlens_k_new=None, "
         "Tensor? seqused_q=None, Tensor? leftpad_k=None, int? page_size=None, int max_seqlen_k_new=0, "
         "bool is_causal=False, int window_size_left=-1, int window_size_right=-1, "
         "int attention_chunk=0, bool has_softcap=False, int num_splits=0, bool? pack_gqa=None, "
         "int sm_margin=0) -> Tensor"
-    ))
+    )))