From 3c46ff2ac5eb522b2a269df1ed4d74830bda3574 Mon Sep 17 00:00:00 2001 From: Baizhou Zhang Date: Fri, 10 Apr 2026 21:39:20 -0700 Subject: [PATCH] fix: restore CPU flash_attn test to use sgl_kernel directly (#22573) Co-authored-by: Claude Opus 4.6 (1M context) --- test/srt/cpu/test_flash_attn.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/test/srt/cpu/test_flash_attn.py b/test/srt/cpu/test_flash_attn.py index 4e1968fa0..e964636d1 100644 --- a/test/srt/cpu/test_flash_attn.py +++ b/test/srt/cpu/test_flash_attn.py @@ -1,12 +1,14 @@ import unittest +import sgl_kernel # noqa: F401 import torch import torch.nn.functional as F from utils import parametrize, precision -from sglang.jit_kernel.flash_attention import flash_attn_varlen_func from sglang.test.test_utils import CustomTestCase +flash_attn_varlen_func = torch.ops.sgl_kernel.flash_attn_varlen_func + torch.manual_seed(1234)