From 73b8eda103d542399198065e78c9e0cb93412137 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 8 May 2026 19:01:49 +0800 Subject: [PATCH] [diffusion] fix: fix FA3 varlen out argument handling (#24688) --- .github/workflows/pr-test.yml | 2 +- python/sglang/jit_kernel/flash_attention_v3.py | 11 ++++++----- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index 9bcf60cd2..79b90aeb5 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -148,7 +148,7 @@ jobs: - ".github/workflows/pr-test-multimodal-gen.yml" - "python/pyproject.toml" - "python/sglang/multimodal_gen/**/!(*.md|*.ipynb)" - - "python/sglang/jit_kernel/diffusion/**" + - "python/sglang/jit_kernel/**" - "python/sglang/jit_kernel/tests/diffusion/**" - "python/sglang/jit_kernel/benchmark/diffusion/**" - "python/sglang/cli/**" diff --git a/python/sglang/jit_kernel/flash_attention_v3.py b/python/sglang/jit_kernel/flash_attention_v3.py index 26781b062..ae4f05026 100644 --- a/python/sglang/jit_kernel/flash_attention_v3.py +++ b/python/sglang/jit_kernel/flash_attention_v3.py @@ -16,15 +16,15 @@ SGL_FA3_KERNEL_REVISION = "v1" DEFAULT_FA3_KERNEL_LOCKFILE = "kernels.lock" -def _call_fa3_kernel(kernel, *args, out=None): +def _call_fa3_kernel(kernel, *args, out=None, **kwargs): if out is None: - return kernel(*args) + return kernel(*args, **kwargs) try: - return kernel(*args, out=out) + return kernel(*args, **kwargs, out=out) except TypeError as exc: if "unexpected keyword argument 'out'" not in str(exc): raise - return kernel(*args) + return kernel(*args, **kwargs) @cache_once @@ -212,7 +212,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=q, k=k, v=v,