yapf

Signed-off-by: Randall Smith <[email protected]>

Author: rasmith

vllm/attention/ops/triton_flash_attention.py

@@ -26,6 +26,7 @@
 
 SUPPORTED_LAYOUTS = ['thd', 'bhsd', 'bshd']
 
+
 class MetaData:
     cu_seqlens_q = None
     cu_seqlens_k = None
@@ -261,7 +262,7 @@ def _attn_fwd_inner(
         # We start from end of seqlen_k so only the first iteration would need
         # to be checked for padding if it is not a multiple of block_n
         # TODO: This can be optimized to only be true for the padded block.
-        if MASK_STEPS: # noqa SIM102
+        if MASK_STEPS:  # noqa SIM102
             # If this is the last block / iteration, we want to
             # mask if the sequence length is not a multiple of block size
             # a solution is to always do BLOCK_M // BLOCK_N + 1 steps if
@@ -633,8 +634,8 @@ def attn_fwd(Q, K, V, bias, SM_SCALE: tl.constexpr, L, Out, stride_qz,
                     # we subtract this from qk which makes it -inf, such that
                     # exp(qk - inf) = 0 for these masked blocks.
                     l_value = tl.full([BLOCK_M],
-                                value=float("inf"),
-                                dtype=tl.float32)
+                                      value=float("inf"),
+                                      dtype=tl.float32)
                     l_ptrs_mask = offs_m < MAX_SEQLENS_Q
                     tl.store(l_ptrs, l_value, mask=l_ptrs_mask)
                     # TODO: Should dropout and return encoded softmax be
@@ -855,7 +856,7 @@ def attn_fwd(Q, K, V, bias, SM_SCALE: tl.constexpr, L, Out, stride_qz,
                 start_m_idx = start_m * BLOCK_M
                 causal_start_idx = seqlen_q - seqlen_k
                 acc = acc.to(Out.type.element_ty)
-                if IS_CAUSAL: # noqa: SIM102
+                if IS_CAUSAL:  # noqa: SIM102
                     if (causal_start_idx > start_m_idx
                             and causal_start_idx < end_m_idx):
                         out_mask_boundary = tl.full((BLOCK_DMODEL, ),
@@ -1358,6 +1359,7 @@ def _attn_bwd(
 
 SUPPORTED_LAYOUTS = ['thd', 'bhsd', 'bshd']
 
+
 def get_shape_from_layout(q, k, metadata):
     assert metadata.layout in SUPPORTED_LAYOUTS, "Got unsupported layout."
     if metadata.layout == 'thd':