fix sgl-kernel unit tests (#5666)

2025-04-23 01:18:30 -07:00
parent e62c49557d
commit 15fabcc07f
9 changed files with 313 additions and 0 deletions
--- a/sgl-kernel/include/sgl_kernel_ops.h
+++ b/sgl-kernel/include/sgl_kernel_ops.h
@@ -352,3 +352,8 @@ std::vector<at::Tensor> mha_varlen_fwd_sparse(
    const bool return_softmax,
    c10::optional<at::Generator> gen_);
 }  // namespace flash
+
+/*
+ * From XGrammar
+ */
+void ApplyTokenBitmaskInplace(at::Tensor logits, at::Tensor bitmask, at::optional<at::Tensor> indices = at::nullopt);