From 03849496adfab14bb3c37bbe7e777b2f260aecf0 Mon Sep 17 00:00:00 2001 From: Mick Date: Sat, 25 Apr 2026 23:42:33 +0800 Subject: [PATCH] jit_kernel: tolerate FA3 kernels without out arg (#23717) --- python/sglang/jit_kernel/flash_attention_v3.py | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/python/sglang/jit_kernel/flash_attention_v3.py b/python/sglang/jit_kernel/flash_attention_v3.py index 732e789c5..00524d5b4 100644 --- a/python/sglang/jit_kernel/flash_attention_v3.py +++ b/python/sglang/jit_kernel/flash_attention_v3.py @@ -15,6 +15,17 @@ SGL_FA3_KERNEL_REVISION = "v1" DEFAULT_FA3_KERNEL_LOCKFILE = "kernels.lock" +def _call_fa3_kernel(kernel, *args, out=None): + if out is None: + return kernel(*args) + try: + return kernel(*args, out=out) + except TypeError as exc: + if "unexpected keyword argument 'out'" not in str(exc): + raise + return kernel(*args) + + @cache_once def _load_fa3_kernels(): # By default, we use the implementation from sgl-kernel, @@ -129,7 +140,8 @@ def flash_attn_with_kvcache( assert k_cache.stride(-1) == 1, "k_cache must have contiguous last dimension" assert v_cache.stride(-1) == 1, "v_cache must have contiguous last dimension" - return _load_fa3_kernels()["flash_attn_with_kvcache"]( + return _call_fa3_kernel( + _load_fa3_kernels()["flash_attn_with_kvcache"], q, k_cache, v_cache, @@ -199,7 +211,8 @@ def flash_attn_varlen_func( "flash_attn at sgl-kernel is only supported on sm90 and above" ) - return _load_fa3_kernels()["flash_attn_varlen_func"]( + return _call_fa3_kernel( + _load_fa3_kernels()["flash_attn_varlen_func"], q, k, v,