From abf69a5e53729c3c75f245f5a5ec36f0c340ce07 Mon Sep 17 00:00:00 2001
From: haosenwang1018 <haosenwang1018@users.noreply.github.com>
Date: Tue, 24 Feb 2026 21:53:03 +0000
Subject: [PATCH] fix: use 'is None' and specific exceptions per PEP 8
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- vibe_attention/interface.py: '== None' → 'is None' for singleton check
- loss/benchmark.py: bare except → except Exception
- sampling/benchmark_all_topk.py: bare except → except Exception
---
 vibe_kernels/loss/benchmark.py              | 6 +++---
 vibe_kernels/sampling/benchmark_all_topk.py | 2 +-
 vibe_kernels/vibe_attention/interface.py    | 2 +-
 3 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/vibe_kernels/loss/benchmark.py b/vibe_kernels/loss/benchmark.py
index 5b483c1..947c7fc 100644
--- a/vibe_kernels/loss/benchmark.py
+++ b/vibe_kernels/loss/benchmark.py
@@ -92,11 +92,11 @@ def cutedsl_impl():
         baseline()
         try:
             triton_impl()
-        except:
+        except Exception:
             pass
         try:
             cutedsl_impl()
-        except:
+        except Exception:
             pass
 
     # --- Quack ---
@@ -177,7 +177,7 @@ def quack_impl():
                 .view(-1)
             )
             tri_diff = torch.max(torch.abs(ref - ours_tri)).item()
-        except:
+        except Exception:
             tri_diff = -1.0
 
         # CuTeDSL
diff --git a/vibe_kernels/sampling/benchmark_all_topk.py b/vibe_kernels/sampling/benchmark_all_topk.py
index 9440f3d..b40b65d 100755
--- a/vibe_kernels/sampling/benchmark_all_topk.py
+++ b/vibe_kernels/sampling/benchmark_all_topk.py
@@ -110,7 +110,7 @@ def benchmark_quack(
             else:
                 # Cutlass doesn't have cuda module in this version
                 return None
-        except:
+        except Exception:
             return None
 
         from quack.topk import topk
diff --git a/vibe_kernels/vibe_attention/interface.py b/vibe_kernels/vibe_attention/interface.py
index b5ac717..930b0c3 100644
--- a/vibe_kernels/vibe_attention/interface.py
+++ b/vibe_kernels/vibe_attention/interface.py
@@ -1276,7 +1276,7 @@ def backward(ctx, dout, *args):
         q, k, v, out, lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k = (
             ctx.saved_tensors
         )
-        assert seqused_q == seqused_k == None
+        assert seqused_q is None and seqused_k is None
         assert ctx.softcap == 0.0
         dq, dk, dv = _flash_attn_bwd(
             q,