[WIP] Add sliding-window attention support to the varlen kernel

build.toml

@@ -1,6 +1,6 @@
 [general]
 name = "metal_flash_sdpa"
-universal = false
+backends = ["metal"]
 
 [torch]
 src = [
@@ -10,9 +10,9 @@ src = [
 
 [kernel.sdpa_metal]
 backend = "metal"
+depends = ["torch"]
 src = [
-  "sdpa-metal/scaled_dot_product_attention.mm",
-  "sdpa-metal/scaled_dot_product_attention.metal",
-  "sdpa-metal/common.h",
+    "sdpa-metal/scaled_dot_product_attention.mm",
+    "sdpa-metal/scaled_dot_product_attention.metal",
+    "sdpa-metal/common.h",
 ]
-depends = [ "torch" ]

sdpa-metal/scaled_dot_product_attention.metal

@@ -1506,6 +1506,10 @@ struct AttnParams {
   int total_k_tokens; ///< Total number of key/value tokens
   int max_seqlen_q; ///< Maximum query sequence length
   int max_seqlen_k; ///< Maximum key/value sequence length
+
+  // Sliding-window attention support (-1 on a side = unbounded on that side)
+  int window_left;  ///< Max distance into the past a query may attend
+  int window_right; ///< Max distance into the future a query may attend
 };
 
 struct AttnMaskParams {
@@ -1521,6 +1525,8 @@ constant bool align_K [[function_constant(201)]];
 
 constant bool has_mask [[function_constant(300)]];
 constant bool do_causal [[function_constant(301)]];
+constant bool has_window [[function_constant(302)]];
+constant bool has_sink [[function_constant(303)]];
 
 template <typename T>
 struct TransformScale {
@@ -1594,6 +1600,7 @@ template <
     const device MaskType* mask [[buffer(6), function_constant(has_mask)]],
     const device int* cu_seqlens_q [[buffer(7)]],  // Cumulative query sequence lengths
     const device int* cu_seqlens_k [[buffer(8)]],  // Cumulative key sequence lengths
+    const device float* sinks [[buffer(9), function_constant(has_sink)]],  // Per-head attention-sink logits
     uint simd_lane_id [[thread_index_in_simdgroup]],
     uint simd_group_id [[simdgroup_index_in_threadgroup]],
     uint3 tid [[threadgroup_position_in_grid]],
@@ -1810,8 +1817,33 @@ template <
     kb_lim = min(kb_lim, (q_block_end_in_seq + BK - 1) / BK);
   }
 
+  // Restrict the key-block range to those intersecting the sliding-window band
+  // for this query block. This is both the perf win (banded, not O(n^2)) and
+  // avoids visiting blocks that are entirely out of band.
+  int kb_start = 0;
+  if (has_window) {
+    int win_offset = (q_seq_len < k_seq_len) ? (k_seq_len - q_seq_len) : 0;
+    int qs = block_idx * BQ + win_offset;          // first query row (key coords)
+    int qe = qs + q_block_size;                    // exclusive
+    if (params->window_left >= 0) {
+      int lo = qs - params->window_left;
+      kb_start = lo > 0 ? lo / BK : 0;
+    }
+    if (params->window_right >= 0) {
+      kb_lim = min(kb_lim, ((qe - 1 + params->window_right) / BK) + 1);
+    }
+  }
+
+  // Fast-forward the K/V block loaders to the first in-band block. The loaders
+  // start at block 0 and only advance via next(), so skipping iterations without
+  // advancing them would read the wrong blocks.
+  for (int s = 0; s < kb_start; ++s) {
+    loader_k.next();
+    loader_v.next();
+  }
+
   // Loop over KV seq length
-  for (int kb = 0; kb < kb_lim; kb++) {
+  for (int kb = kb_start; kb < kb_lim; kb++) {
     // Load K block and apply scale
     threadgroup_barrier(mem_flags::mem_threadgroup);
     
@@ -1894,6 +1926,38 @@ template <
       }
     }
 
+    // Mask out keys outside the sliding window band [row - window_left, row + window_right]
+    if (has_window) {
+      using stile_t = decltype(Stile);
+      using selem_t = typename stile_t::elem_type;
+      constexpr auto neg_inf = -metal::numeric_limits<selem_t>::infinity();
+
+      const int wl = params->window_left;   // -1 => unbounded into the past
+      const int wr = params->window_right;  // -1 => unbounded into the future
+
+      STEEL_PRAGMA_UNROLL
+      for (short i = 0; i < stile_t::kTileRows; i++) {
+        // Same row-position machinery as the causal block above.
+        int row_pos = block_idx * BQ + tm + sm + (i * stile_t::kFragRows);
+        if (q_seq_len < k_seq_len) {
+          row_pos += (k_seq_len - q_seq_len);
+        }
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < stile_t::kTileCols; j++) {
+          const int col_pos_in_seq = kb * BK + sn + (j * stile_t::kFragCols);
+          STEEL_PRAGMA_UNROLL
+          for (short jj = 0; jj < stile_t::MMAFrag_t::kElemCols; jj++) {
+            const int col = col_pos_in_seq + jj;
+            const bool past_ok = (wl < 0) || ((row_pos - col) <= wl);
+            const bool future_ok = (wr < 0) || ((col - row_pos) <= wr);
+            if (!(past_ok && future_ok)) {
+              Stile.frag_at(i, j)[jj] = neg_inf;
+            }
+          }
+        }
+      }
+    }
+
     // Other masking as needed
     if (has_mask) {
       using stile_t = decltype(Stile);
@@ -1979,6 +2043,18 @@ template <
     // Row max
     Stile.template row_reduce<MaxOp>(new_max);
 
+    // A sliding-window query row can have a key block fully masked to -inf.
+    // Then new_max stays -inf and exp2(-inf - (-inf)) = NaN poisons the row.
+    // Replace a -inf running max with a finite value (the prior max if any,
+    // else 0): masked entries then exp2 to 0 and accumulators stay 0.
+    constexpr AccumType row_ninf = -metal::numeric_limits<AccumType>::infinity();
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kRowsPT; ++i) {
+      if (new_max[i] == row_ninf) {
+        new_max[i] = (max_score[i] == row_ninf) ? AccumType(0) : max_score[i];
+      }
+    }
+
     // exp(Si - rowmax(Si))
     Stile.template row_bin_op<ExpSubOp>(new_max);
 
@@ -2044,6 +2120,24 @@ template <
     loader_v.next();
   }
 
+  // Attention sink: an extra per-head logit that only enters the softmax
+  // denominator (its "value" is zero), matching gpt-oss s_aux. Scores live in
+  // the base-2 domain (scale folds in log2(e)), so convert the raw sink the
+  // same way and re-fold the running max for numerical stability.
+  if (has_sink) {
+    const AccumType sink_log2 = static_cast<AccumType>(sinks[head_idx]) * 1.44269504089f;
+    AccumType sink_factor[kRowsPT];
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kRowsPT; ++i) {
+      const AccumType m = metal::max(max_score[i], sink_log2);
+      const AccumType resc = fast::exp2(max_score[i] - m);
+      sum_score[i] = sum_score[i] * resc + fast::exp2(sink_log2 - m);
+      sink_factor[i] = resc;
+      max_score[i] = m;
+    }
+    Otile.template row_bin_op<MulOp>(sink_factor);
+  }
+
   // Normalize output
   Otile.template row_bin_op<DivOp>(sum_score);
   threadgroup_barrier(mem_flags::mem_none);

sdpa-metal/scaled_dot_product_attention.mm

@@ -69,6 +69,8 @@ struct AttnParams {
   int32_t total_k_tokens; // Total number of key/value tokens
   int32_t max_seqlen_q;   // Maximum query sequence length
   int32_t max_seqlen_k;   // Maximum key/value sequence length
+  int32_t window_left;    // Sliding window: max distance into the past (-1 = unbounded)
+  int32_t window_right;   // Sliding window: max distance into the future (-1 = unbounded)
 };
 
 // Forward declarations for kernel implementations
@@ -86,7 +88,10 @@ void call_flash_attention_varlen(
     int64_t max_seqlen_k,
     bool do_causal,
     double scale,
-    double softcapping);
+    double softcapping,
+    int64_t window_left,
+    int64_t window_right,
+    const std::optional<torch::Tensor> &sinks);
 
 
 void flash_attention_varlen(
@@ -100,7 +105,10 @@ void flash_attention_varlen(
     int64_t max_seqlen_k,         // Maximum key sequence length
     bool do_causal,               // Whether to use causal mask
     double scale,                 // Attention scale
-    double softcapping) {         // Softcapping value
+    double softcapping,           // Softcapping value
+    int64_t window_left,          // Sliding window past extent (-1 = unbounded)
+    int64_t window_right,         // Sliding window future extent (-1 = unbounded)
+    const std::optional<torch::Tensor> &sinks) {  // Per-head attention-sink logits (fp32)
 
   try {
     // Get device and stream
@@ -142,9 +150,9 @@ void flash_attention_varlen(
   // For variable-length Flash Attention, always use the full attention kernel
   
   // Call the Flash Attention kernel
-  call_flash_attention_varlen(device, cmdBuf, lib, out, query, key, value, 
+  call_flash_attention_varlen(device, cmdBuf, lib, out, query, key, value,
                               cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k,
-                              do_causal, scale, softcapping);
+                              do_causal, scale, softcapping, window_left, window_right, sinks);
   } catch (const std::exception& e) {
     throw;
   } catch (...) {
@@ -167,7 +175,10 @@ void call_flash_attention_varlen(
     int64_t max_seqlen_k,
     bool do_causal,
     double scale,
-    double softcapping) {
+    double softcapping,
+    int64_t window_left,
+    int64_t window_right,
+    const std::optional<torch::Tensor> &sinks) {
 
   // Get dimensions
   int64_t total_q_tokens = query.size(0);
@@ -197,7 +208,9 @@ void call_flash_attention_varlen(
   params.total_k_tokens = key.size(0);
   params.max_seqlen_q = max_seqlen_q;
   params.max_seqlen_k = max_seqlen_k;
-  
+  params.window_left = static_cast<int32_t>(window_left);
+  params.window_right = static_cast<int32_t>(window_right);
+
   // Initialize fields that might be checked but aren't used in Flash Attention
   params.qL = 0;  // Not used in variable-length attention
   params.kL = 0;  // Not used in variable-length attention
@@ -227,11 +240,15 @@ void call_flash_attention_varlen(
   // The kernel will handle the cu_seqlens internally
 
   bool has_mask = false;  // Masks are not supported in Flash Attention
+  bool has_window = (window_left >= 0) || (window_right >= 0);
+  bool has_sink = sinks.has_value();
 
   // Setup function constants
   MTLFunctionConstantValues *constants = [MTLFunctionConstantValues new];
   [constants setConstantValue:&has_mask type:MTLDataTypeBool atIndex:300];
   [constants setConstantValue:&do_causal type:MTLDataTypeBool atIndex:301];
+  [constants setConstantValue:&has_window type:MTLDataTypeBool atIndex:302];
+  [constants setConstantValue:&has_sink type:MTLDataTypeBool atIndex:303];
 
   // Construct kernel name based on data type and head dimension
   std::string kernel_name = "steel_attention_";
@@ -259,6 +276,10 @@ void call_flash_attention_varlen(
   at::mps::MPSStream *stream = at::mps::getCurrentMPSStream();
   dispatch_queue_t q = stream->queue();
   dispatch_sync(q, ^{
+    // The MPS stream may already have an open compute encoder (e.g. coalesced
+    // kernels under torch.compile's inductor backend). Close it before opening
+    // ours, otherwise Metal asserts "A command encoder is already encoding".
+    stream->endKernelCoalescing();
     id<MTLComputeCommandEncoder> encoder = [cmdBuf computeCommandEncoder];
     TORCH_CHECK(encoder, "Failed to create compute encoder");
 
@@ -297,10 +318,18 @@ void call_flash_attention_varlen(
     [encoder setBuffer:getMTLBufferStorage(cu_seqlens_q) 
               offset:cu_seqlens_q.storage_offset() * cu_seqlens_q.element_size() 
               atIndex:7];
-    [encoder setBuffer:getMTLBufferStorage(cu_seqlens_k) 
-              offset:cu_seqlens_k.storage_offset() * cu_seqlens_k.element_size() 
+    [encoder setBuffer:getMTLBufferStorage(cu_seqlens_k)
+              offset:cu_seqlens_k.storage_offset() * cu_seqlens_k.element_size()
               atIndex:8];
 
+    // Per-head attention-sink logits - index 9 (only bound when has_sink)
+    if (has_sink) {
+      const torch::Tensor &sinks_t = sinks.value();
+      [encoder setBuffer:getMTLBufferStorage(sinks_t)
+                offset:sinks_t.storage_offset() * sinks_t.element_size()
+                atIndex:9];
+    }
+
     // Calculate grid dimensions
     // We need to process each sequence independently
     int64_t max_blocks_q = (max_seqlen_q + BQ - 1) / BQ;

torch-ext/metal_flash_sdpa/__init__.py

@@ -1,11 +1,13 @@
 from ._custom_ops import (
     flash_attention_varlen,
+    flash_attn_func,
     flash_attn_varlen_func,
 )
 from ._ops import ops
 
 __all__ = [
     "flash_attention_varlen",
+    "flash_attn_func",
     "flash_attn_varlen_func",
     "ops",
 ]

torch-ext/metal_flash_sdpa/_custom_ops.py

@@ -17,6 +17,9 @@ def flash_attention_varlen(
     do_causal: bool = False,
     scale: Optional[float] = None,
     softcapping: float = 1.0,
+    window_left: int = -1,
+    window_right: int = -1,
+    sinks: Optional[torch.Tensor] = None,
 ) -> None:
     """
     Flash Attention with variable-length sequences.
@@ -38,10 +41,20 @@ def flash_attention_varlen(
         - cu_seqlens_q and cu_seqlens_k must have dtype torch.int32 for Metal compatibility
         - Supported head dimensions: 32, 64, 72, 80, 96, 128
         - Masks are not supported
+        - window_left / window_right bound a sliding-window band (-1 = unbounded)
     """
     if scale is None:
         scale = query.shape[-1] ** -0.5
-    
+
+    # The kernel reads raw storage assuming contiguous [tokens, heads, dim]. Callers may pass
+    # transposed/reshaped views, so force contiguity on the inputs (out must already be contiguous).
+    query, key, value = query.contiguous(), key.contiguous(), value.contiguous()
+
+    # The kernel reads `sinks` as fp32 (the model keeps them in fp32). transformers' flash
+    # wrapper casts s_aux to the query dtype (fp16), so cast back here to avoid misreads.
+    if sinks is not None:
+        sinks = sinks.to(torch.float32).contiguous()
+
     ops.flash_attention_varlen(
         out,
         query,
@@ -54,6 +67,9 @@ def flash_attention_varlen(
         do_causal,
         scale,
         softcapping,
+        window_left,
+        window_right,
+        sinks,
     )
 
 def flash_attn_varlen_func(
@@ -71,25 +87,29 @@ def flash_attn_varlen_func(
     alibi_slopes: Optional[torch.Tensor] = None,
     deterministic: bool = False,
     return_attn_probs: bool = False,
+    s_aux: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
     """
     Flash Attention function with API compatible with the original Flash Attention.
     
     Note: This implementation does not support:
     - dropout
-    - window attention
     - alibi slopes
     - returning attention probabilities
+
+    `window_size = (left, right)` follows the flash-attn convention: a token attends to
+    keys in [pos - left, pos + right]; -1 means unbounded on that side.
     """
     if dropout_p > 0:
         raise NotImplementedError("Dropout is not supported in this implementation")
-    if window_size != (-1, -1):
-        raise NotImplementedError("Window attention is not supported")
     if alibi_slopes is not None:
         raise NotImplementedError("ALiBi is not supported")
     if return_attn_probs:
         raise NotImplementedError("Returning attention probabilities is not supported")
-    
+
+    # Ensure contiguous so the output buffer (empty_like) and inputs match the kernel's layout.
+    q, k, v = q.contiguous(), k.contiguous(), v.contiguous()
+
     # Create output tensor
     out = torch.empty_like(q)
     
@@ -106,12 +126,63 @@ def flash_attn_varlen_func(
         do_causal=causal,
         scale=softmax_scale,
         softcapping=1.0,
+        window_left=window_size[0],
+        window_right=window_size[1],
+        sinks=s_aux,
     )
-    
+
     return out
 
 
+def flash_attn_func(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    dropout_p: float = 0.0,
+    softmax_scale: Optional[float] = None,
+    causal: bool = False,
+    window_size: tuple = (-1, -1),
+    alibi_slopes: Optional[torch.Tensor] = None,
+    deterministic: bool = False,
+    return_attn_probs: bool = False,
+    s_aux: Optional[torch.Tensor] = None,
+    **kwargs,
+) -> torch.Tensor:
+    """
+    Non-varlen entry point: q/k/v are [batch, seqlen, heads, head_dim] with uniform length.
+    Wraps the varlen kernel by treating each batch row as one packed sequence. Lets the kernel
+    be used from the standard (non-continuous-batching) attention path.
+    """
+    if dropout_p and dropout_p > 0:
+        raise NotImplementedError("Dropout is not supported in this implementation")
+    if alibi_slopes is not None:
+        raise NotImplementedError("ALiBi is not supported")
+    if return_attn_probs:
+        raise NotImplementedError("Returning attention probabilities is not supported")
+
+    B, S, H, D = q.shape
+    Sk = k.shape[1]
+    cu_q = torch.arange(0, (B + 1) * S, S, device=q.device, dtype=torch.int32)
+    cu_k = torch.arange(0, (B + 1) * Sk, Sk, device=k.device, dtype=torch.int32)
+    out = flash_attn_varlen_func(
+        q.reshape(B * S, H, D),
+        k.reshape(B * Sk, k.shape[2], D),
+        v.reshape(B * Sk, v.shape[2], D),
+        cu_q,
+        cu_k,
+        S,
+        Sk,
+        dropout_p=0.0,
+        softmax_scale=softmax_scale,
+        causal=causal,
+        window_size=window_size,
+        s_aux=s_aux,
+    )
+    return out.reshape(B, S, H, D)
+
+
 __all__ = [
     "flash_attention_varlen",
     "flash_attn_varlen_func",
+    "flash_attn_func",
 ]

torch-ext/torch_binding.cpp

@@ -4,7 +4,7 @@
 #include "torch_binding.h"
 
 TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
-  ops.def("flash_attention_varlen(Tensor! out, Tensor query, Tensor key, Tensor value, Tensor cu_seqlens_q, Tensor cu_seqlens_k, int max_seqlen_q, int max_seqlen_k, bool do_causal, float scale, float softcapping) -> ()");
+  ops.def("flash_attention_varlen(Tensor! out, Tensor query, Tensor key, Tensor value, Tensor cu_seqlens_q, Tensor cu_seqlens_k, int max_seqlen_q, int max_seqlen_k, bool do_causal, float scale, float softcapping, int window_left, int window_right, Tensor? sinks) -> ()");
   ops.impl("flash_attention_varlen", torch::kMPS, flash_attention_varlen);
 }
 

torch-ext/torch_binding.h

@@ -13,4 +13,7 @@ void flash_attention_varlen(
     int64_t max_seqlen_k,
     bool do_causal,
     double scale,
-    double softcapping);
+    double softcapping,
+    int64_t window_left,
+    int64_t window_right,
+    const std::optional<torch::Tensor> &sinks);