feat: integrate xqa attention backend (#1503)

<!-- .github/pull_request_template.md -->

## 📌 Description

<!-- What does this PR do? Briefly describe the changes and why they’re
needed. -->

## 🔍 Related Issues

<!-- Link any related issues here -->

## 🚀 Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull
request, please make sure the following items are complete.

### ✅ Pre-commit Checks

- [x] I have installed `pre-commit` by running `pip install pre-commit`
(or used your preferred method).
- [x] I have installed the hooks with `pre-commit install`.
- [x] I have run the hooks manually with `pre-commit run --all-files`
and fixed any reported issues.

> If you are unsure about how to set up `pre-commit`, see [the
pre-commit documentation](https://pre-commit.com/).

## 🧪 Tests

- [ ] Tests have been added or updated as needed.
- [ ] All tests are passing (`unittest`, etc.).

## Reviewer Notes

<!-- Optional: anything you'd like reviewers to focus on, concerns, etc.
-->

---------

Signed-off-by: Qidi Sang <[email protected]>

Author: qsang-nv

csrc/flashinfer_xqa_ops.cu

@@ -0,0 +1,35 @@
+/*
+ * Copyright (c) 2024 by FlashInfer team.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "pytorch_extension_utils.h"
+
+void xqa_wrapper(int64_t multiProcessorCount, int64_t nbKHeads, int64_t slidingWinSize,
+                 double qScale, at::Tensor output,
+#if LOW_PREC_OUTPUT
+                 at::Tensor rcpOutScale,
+#endif
+                 at::Tensor q, at::Tensor attentionSinks, at::Tensor pool,
+                 at::Tensor kvCachePageList, int64_t maxSeqLen, at::Tensor seqLen,
+                 int64_t batchSize, at::Tensor kvCacheScale,
+#if SPEC_DEC
+                 int64_t qSeqLen, at::Tensor qCuSeqLens, at::Tensor mask,
+#endif
+                 at::Tensor semaphores, at::Tensor scratch);
+
+TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) {
+  // "XQA Wrapper"
+  m.def("xqa_wrapper", xqa_wrapper);
+}