From 93590eb1adfc6ca345748c04783c2e82310de39d Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sun, 24 May 2026 03:58:12 +0000 Subject: [PATCH] D1: Fix syntax (separate kv_stage line) --- dsv4/kernels/attention/fmha.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/dsv4/kernels/attention/fmha.py b/dsv4/kernels/attention/fmha.py index f700f844..722527fb 100644 --- a/dsv4/kernels/attention/fmha.py +++ b/dsv4/kernels/attention/fmha.py @@ -30,7 +30,8 @@ class FmhaKernel: self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE self.epilogue_warp_id = (0,1,2,3); self.mma_warp_id = 4; self.tma_warp_id = 5 self.threads_per_cta = 192; self.num_c_stage = 2 - self.kv_stage = 1 if head_dim > 128 else 2 # Reduce SMEM at large hd; self.q_stage = 1; self.num_c_stage = 2 + self.kv_stage = 1 if head_dim > 128 else 2 # Reduce SMEM at large hd + self.q_stage = 1; self.num_c_stage = 2 self.scale_softmax = scale_softmax if scale_softmax is not None else 1.0 / math.sqrt(self.head_dim) self.scale_softmax_log2 = self.scale_softmax * math.log2(math.e)