feat(flash_attention): add vLLM-style sliding_window entry (additive)

Keeps the native `window_left` / `window_right` pair as-is and adds an
optional `std::optional<int64_t> sliding_window` parameter.  When set,
the base class normalizes it to the causal-sliding pair
`(sliding_window - 1, 0)`; when both forms are supplied the normalized
values must agree.  Callers can now use either entry point:

  // Pair form (existing, unchanged):
  flash_attention(..., window_left=255, window_right=0, ...)

  // vLLM form:
  flash_attention(..., sliding_window=256, ...)

Ascend impl reads the resolved pair from the base-class members
(`window_left_` / `window_right_`) so `sliding_window` is honored at
both construction and call time.

Also extends `generate_wrappers.py` to set `py::arg(...) = py::none()`
defaults for all `std::optional<...>` parameters (previously only
`std::optional<Tensor>`), so `sliding_window` is properly optional on
the Python side.

Adds `test_flash_attention_sliding_window_equivalence` asserting
bit-exact equality between the two entry points.

Author: zhangyue

scripts/generate_wrappers.py

@@ -121,6 +121,9 @@ def _is_optional_tensor(arg):
             return True
         return "std::optional" in arg.type.spelling and "Tensor" in arg.type.spelling
 
+    def _is_optional(arg):
+        return "std::optional" in arg.type.spelling
+
     def _is_vector_tensor(arg):
         if arg.spelling in vector_tensor_params:
             return True
@@ -177,7 +180,7 @@ def _generate_py_args(node):
             if arg.spelling == "stream":
                 continue
 
-            if _is_optional_tensor(arg):
+            if _is_optional(arg):
                 parts.append(f'py::arg("{arg.spelling}") = py::none()')
             else:
                 parts.append(f'py::arg("{arg.spelling}")')

src/ascend/flash_attention/kernel.h

@@ -114,10 +114,12 @@ class Operator<FlashAttention, Device::Type::kAscend> : public FlashAttention {
            std::optional<Tensor> block_table, int64_t num_heads,
            int64_t num_kv_heads, int64_t head_size, double scale, bool causal,
            int64_t window_left, int64_t window_right, int64_t block_size,
-           Tensor output)
+           Tensor output,
+           std::optional<int64_t> sliding_window = std::nullopt)
       : FlashAttention(query, key, value, cu_seqlens_q, cu_seqlens_kv,
                        block_table, num_heads, num_kv_heads, head_size, scale,
-                       causal, window_left, window_right, block_size, output) {
+                       causal, window_left, window_right, block_size, output,
+                       sliding_window) {
     paged_ = block_table.has_value() && block_size > 0;
     aclDataType acl_dt = ascend::toAclDtype(query.dtype());
 
@@ -126,9 +128,11 @@ class Operator<FlashAttention, Device::Type::kAscend> : public FlashAttention {
       prefill_q_cache_ = ascend::AclTensorCache(query);
       prefill_out_cache_ = ascend::AclTensorCache(output);
 
-      // Pre-compute causal mask once (sparse_mode >= 2).
+      // Pre-compute causal mask once (sparse_mode >= 2).  Read the
+      // resolved pair from base-class members so `sliding_window`
+      // normalization is honored at cache-key construction.
       if (causal) {
-        int64_t sm = (window_left >= 0) ? 4 : 3;
+        int64_t sm = (window_left_ >= 0) ? 4 : 3;
         if (sm >= 2) {
           causal_mask_ = detail::makeCausalMask(&causal_mask_buf_, nullptr);
         }
@@ -169,26 +173,36 @@ class Operator<FlashAttention, Device::Type::kAscend> : public FlashAttention {
                   std::optional<Tensor> block_table, int64_t num_heads,
                   int64_t num_kv_heads, int64_t head_size, double scale,
                   bool causal, int64_t window_left, int64_t window_right,
-                  int64_t block_size, Tensor output) const override {
+                  int64_t block_size, Tensor output,
+                  std::optional<int64_t> sliding_window) const override {
     auto stream = static_cast<aclrtStream>(stream_);
     const bool paged = paged_;
 
+    // The base class stored the resolved window pair in `window_left_` /
+    // `window_right_` at construction; prefer those over the call-site
+    // args so that `sliding_window` is honored here as well.
+    int64_t wl = window_left_;
+    int64_t wr = window_right_;
+    (void)window_left;
+    (void)window_right;
+    (void)sliding_window;
+
     int64_t sparse_mode;
     int64_t pre_tokens = 2147483647;
     int64_t next_tokens = 2147483647;
     if (causal) {
-      if (window_left >= 0) {
+      if (wl >= 0) {
         sparse_mode = 4;
-        pre_tokens = window_left;
+        pre_tokens = wl;
         next_tokens = 0;
       } else {
         sparse_mode = 3;
         next_tokens = 0;
       }
     } else {
       sparse_mode = 0;
-      if (window_left >= 0) pre_tokens = window_left;
-      if (window_right >= 0) next_tokens = window_right;
+      if (wl >= 0) pre_tokens = wl;
+      if (wr >= 0) next_tokens = wr;
     }
 
     if (!paged) {

src/base/flash_attention.h

@@ -11,21 +11,30 @@ namespace infini::ops {
 
 class FlashAttention : public Operator<FlashAttention> {
  public:
+  // `window_left` / `window_right` is the native InfiniOps pair-form
+  // window (left-context / right-context tokens, `-1` = disabled).
+  // `sliding_window` is a vLLM-style single-parameter shortcut: when
+  // set, it is normalized to `(sliding_window - 1, 0)` — i.e. causal
+  // sliding over the most recent `sliding_window` tokens.  When both
+  // forms are supplied the normalized values must agree.  Callers may
+  // use whichever form is more natural; the kernel only sees the
+  // resolved pair.
   FlashAttention(const Tensor query, const Tensor key, const Tensor value,
                  std::optional<Tensor> cu_seqlens_q,
                  std::optional<Tensor> cu_seqlens_kv,
                  std::optional<Tensor> block_table, int64_t num_heads,
                  int64_t num_kv_heads, int64_t head_size, double scale,
                  bool causal, int64_t window_left, int64_t window_right,
-                 int64_t block_size, Tensor output)
+                 int64_t block_size, Tensor output,
+                 std::optional<int64_t> sliding_window = std::nullopt)
       : num_tokens_{query.size(0)},
         num_heads_{num_heads},
         num_kv_heads_{num_kv_heads},
         head_size_{head_size},
         scale_{scale},
         causal_{causal},
-        window_left_{window_left},
-        window_right_{window_right},
+        window_left_{resolveWindowLeft(window_left, sliding_window)},
+        window_right_{resolveWindowRight(window_right, sliding_window)},
         block_size_{block_size},
         dtype_{query.dtype()},
         query_shape_{query.shape()},
@@ -45,15 +54,37 @@ class FlashAttention : public Operator<FlashAttention> {
            "`FlashAttention` requires query to be 3D [T, N, D]");
   }
 
-  virtual void operator()(const Tensor query, const Tensor key,
-                          const Tensor value,
-                          std::optional<Tensor> cu_seqlens_q,
-                          std::optional<Tensor> cu_seqlens_kv,
-                          std::optional<Tensor> block_table, int64_t num_heads,
-                          int64_t num_kv_heads, int64_t head_size, double scale,
-                          bool causal, int64_t window_left,
-                          int64_t window_right, int64_t block_size,
-                          Tensor output) const = 0;
+  virtual void operator()(
+      const Tensor query, const Tensor key, const Tensor value,
+      std::optional<Tensor> cu_seqlens_q, std::optional<Tensor> cu_seqlens_kv,
+      std::optional<Tensor> block_table, int64_t num_heads, int64_t num_kv_heads,
+      int64_t head_size, double scale, bool causal, int64_t window_left,
+      int64_t window_right, int64_t block_size, Tensor output,
+      std::optional<int64_t> sliding_window = std::nullopt) const = 0;
+
+ private:
+  // Normalize the window representation.  If both the explicit pair and
+  // `sliding_window` are supplied, assert the pair matches the derived
+  // `(sliding_window - 1, 0)` causal-sliding window.
+  static int64_t resolveWindowLeft(int64_t window_left,
+                                   std::optional<int64_t> sliding_window) {
+    if (!sliding_window.has_value()) return window_left;
+    int64_t derived = sliding_window.value() - 1;
+    assert((window_left == -1 || window_left == derived) &&
+           "`FlashAttention`: `window_left` inconsistent with `sliding_window`");
+    return derived;
+  }
+
+  static int64_t resolveWindowRight(int64_t window_right,
+                                    std::optional<int64_t> sliding_window) {
+    if (!sliding_window.has_value()) return window_right;
+    assert((window_right == -1 || window_right == 0) &&
+           "`FlashAttention`: `window_right` inconsistent with `sliding_window` "
+           "(vLLM sliding_window implies right=0)");
+    return 0;
+  }
+
+ public:
 
  protected:
   Tensor::Size num_tokens_{0};

tests/test_flash_attention.py

@@ -537,3 +537,80 @@ def _ref_flash_attention_paged(
         outputs.append(out)
 
     return torch.cat(outputs, dim=0).to(query.device)
+
+
+@pytest.mark.parametrize("sliding_window", (4, 16))
+@pytest.mark.parametrize("device", ("npu",))
+def test_flash_attention_sliding_window_equivalence(sliding_window, device):
+    """The vLLM-style `sliding_window=N` entry must produce the same output
+    as the native `window_left=N-1, window_right=0` pair.
+    """
+    if not (hasattr(torch, "npu") and torch.npu.is_available()):
+        pytest.skip("NPU not available")
+
+    num_tokens = 32
+    num_heads = 8
+    num_kv_heads = 8
+    head_size = 64
+    scale = 1.0 / head_size**0.5
+    dtype = torch.float16
+
+    query = randn_strided(
+        (num_tokens, num_heads, head_size), None, dtype=dtype, device=device
+    )
+    key = randn_strided(
+        (num_tokens, num_kv_heads, head_size), None, dtype=dtype, device=device
+    )
+    value = randn_strided(
+        (num_tokens, num_kv_heads, head_size), None, dtype=dtype, device=device
+    )
+
+    cu_seqlens_q = torch.tensor([0, num_tokens], dtype=torch.int64, device=device)
+    cu_seqlens_kv = torch.tensor([0, num_tokens], dtype=torch.int64, device=device)
+
+    # Pair-form call.
+    out_pair = torch.empty_like(query)
+    infini.ops.flash_attention(
+        query,
+        key,
+        value,
+        cu_seqlens_q,
+        cu_seqlens_kv,
+        None,
+        num_heads,
+        num_kv_heads,
+        head_size,
+        scale,
+        True,
+        sliding_window - 1,
+        0,
+        0,
+        out_pair,
+        stream=get_npu_stream(query),
+    )
+
+    # vLLM-style single-parameter call.
+    out_sw = torch.empty_like(query)
+    infini.ops.flash_attention(
+        query,
+        key,
+        value,
+        cu_seqlens_q,
+        cu_seqlens_kv,
+        None,
+        num_heads,
+        num_kv_heads,
+        head_size,
+        scale,
+        True,
+        -1,
+        -1,
+        0,
+        out_sw,
+        sliding_window=sliding_window,
+        stream=get_npu_stream(query),
+    )
+
+    assert torch.equal(out_pair, out_sw), (
+        f"Max diff: {(out_pair.float() - out_sw.float()).abs().max().item()}"
+    )