lint

Author: Rachmanino

examples/attention_sink/example_gqa_sink_bwd_bhsd.py

@@ -20,8 +20,7 @@ def get_bwd_configs():
 
 
 @tilelang.jit(
-    out_idx=[3, 4],
-    pass_configs={
+    out_idx=[3, 4], pass_configs={
         tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True,
     })
 def flashattn_fwd(
@@ -139,8 +138,7 @@ def flash_fwd(
 
 
 @tilelang.jit(
-    out_idx=[2],
-    pass_configs={
+    out_idx=[2], pass_configs={
         tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True,
     })
 def flashattn_bwd_preprocess(batch, heads, seq_len, dim, dtype: str = "float16"):
@@ -178,8 +176,7 @@ def make_dq_layout(dQ):
 
 
 @tilelang.jit(
-    out_idx=[1],
-    pass_configs={
+    out_idx=[1], pass_configs={
         tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True,
     })
 def flashattn_bwd_postprocess(batch, heads, seq_len, dim, dtype: str = "float16"):
@@ -202,10 +199,9 @@ def flash_bwd_post(
     return flash_bwd_post
 
 
-@tilelang.jit(
-    pass_configs={
-        tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True,
-    })
+@tilelang.jit(pass_configs={
+    tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True,
+})
 def flashattn_bwd(batch,
                   heads,
                   seq_len,

examples/attention_sink/example_gqa_sink_fwd_bhsd_wgmma_pipelined.py

@@ -23,8 +23,7 @@ def get_configs():
     rep=100,
 )
 @tilelang.jit(
-    out_idx=[3],
-    pass_configs={
+    out_idx=[3], pass_configs={
         tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True,
     })
 def flashattn(

examples/attention_sink/example_mha_sink_bwd_bhsd.py

@@ -20,8 +20,7 @@ def get_bwd_configs():
 
 
 @tilelang.jit(
-    out_idx=[3, 4],
-    pass_configs={
+    out_idx=[3, 4], pass_configs={
         tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True,
     })
 def flashattn_fwd(
@@ -136,8 +135,7 @@ def flash_fwd(
 
 
 @tilelang.jit(
-    out_idx=[2],
-    pass_configs={
+    out_idx=[2], pass_configs={
         tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True,
     })
 def flashattn_bwd_preprocess(batch, heads, seq_len, dim, dtype: str = "float16"):
@@ -175,8 +173,7 @@ def make_dq_layout(dQ):
 
 
 @tilelang.jit(
-    out_idx=[1],
-    pass_configs={
+    out_idx=[1], pass_configs={
         tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True,
     })
 def flashattn_bwd_postprocess(batch, heads, seq_len, dim, dtype: str = "float16"):
@@ -199,10 +196,9 @@ def flash_bwd_post(
     return flash_bwd_post
 
 
-@tilelang.jit(
-    pass_configs={
-        tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True,
-    })
+@tilelang.jit(pass_configs={
+    tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True,
+})
 def flashattn_bwd(
     batch,
     heads,

examples/attention_sink/example_mha_sink_fwd_bhsd.py

@@ -18,8 +18,7 @@ def get_configs():
 
 @autotune(configs=get_configs(), warmup=500, rep=100)
 @tilelang.jit(
-    out_idx=[3],
-    pass_configs={
+    out_idx=[3], pass_configs={
         tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True,
     })
 def flashattn(

examples/attention_sink/example_mha_sink_fwd_bhsd_wgmma_pipelined.py

@@ -19,8 +19,7 @@ def get_configs():
 
 @autotune(configs=get_configs(), warmup=500, rep=100)
 @tilelang.jit(
-    out_idx=[3],
-    pass_configs={
+    out_idx=[3], pass_configs={
         tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True,
     })
 def flashattn(

src/target/codegen_cuda.cc

@@ -948,17 +948,21 @@ void CodeGenTileLangCUDA::VisitExpr_(const CastNode *op, std::ostream &os) {
     if (from_ty.lanes() == 2 && target_ty.lanes() == 2) {
       // bfloat162 -> float2
       PrintIndent();
-      stream << sret << " = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162*>(&(" << src << ")));\n";
+      stream << sret
+             << " = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162*>(&("
+             << src << ")));\n";
       os << sret;
       return;
     } else if (from_ty.lanes() == 4 && target_ty.lanes() == 4) {
       // bfloat162x2 -> float4
       PrintIndent();
       stream << "((float2*)(&" << sret << "))[0] = "
-             << "__bfloat1622float2(*reinterpret_cast<__nv_bfloat162*>(&(" << src << ")));\n";
+             << "__bfloat1622float2(*reinterpret_cast<__nv_bfloat162*>(&("
+             << src << ")));\n";
       PrintIndent();
       stream << "((float2*)(&" << sret << "))[1] = "
-             << "__bfloat1622float2(*(reinterpret_cast<__nv_bfloat162*>(&(" << src << "))+1));\n";
+             << "__bfloat1622float2(*(reinterpret_cast<__nv_bfloat162*>(&("
+             << src << "))+1));\n";
       os << sret;
       return;
     }
@@ -967,8 +971,8 @@ void CodeGenTileLangCUDA::VisitExpr_(const CastNode *op, std::ostream &os) {
     if (from_ty.lanes() == 2 && target_ty.lanes() == 2) {
       // float2 -> bfloat162
       PrintIndent();
-      stream << "*reinterpret_cast<__nv_bfloat162*>(&(" << sret << ")) = __float22bfloat162_rn(*(float2*)(&("
-             << src << ")));\n";
+      stream << "*reinterpret_cast<__nv_bfloat162*>(&(" << sret
+             << ")) = __float22bfloat162_rn(*(float2*)(&(" << src << ")));\n";
       os << sret;
       return;
     } else if (from_ty.lanes() == 4 && target_ty.lanes() == 4) {

testing/python/language/test_tilelang_language_vectorized_cast.py

@@ -32,7 +32,6 @@ def run_vectorized_cast(src_dtype_str: str, dst_dtype_str: str, check_str: str,
         src_dtype_str: The source data type string.
         dst_dtype_str: The destination data type string.
         check_str: Used to ensure vectorized cast is used.
-        M: The size of the tensor.
         lanes: The number of lanes of the source and destination data types.
     """
 
@@ -80,4 +79,3 @@ def test_vectorized_cast():
 
 if __name__ == "__main__":
     tilelang.testing.main()
-