From ebe221e6487894e5b4b8c15d68deb0c364c57678 Mon Sep 17 00:00:00 2001 From: sufubao Date: Wed, 14 Oct 2026 18:55:17 +0800 Subject: [PATCH 1/2] fix: pass required arguments to SGL FA3 vision attention --- .../common/basemodel/attention_vit/fa3/fp.py | 2 ++ .../vit/triton_kernel/flashattention_nopad.py | 2 ++ .../vit/test_flash_attention_forward.py | 30 +++++++++++++++++-- 3 files changed, 31 insertions(+), 3 deletions(-) diff --git a/lightllm/common/basemodel/attention_vit/fa3/fp.py b/lightllm/common/basemodel/attention_vit/fa3/fp.py index f804116f1f..a8e56417c5 100644 --- a/lightllm/common/basemodel/attention_vit/fa3/fp.py +++ b/lightllm/common/basemodel/attention_vit/fa3/fp.py @@ -53,6 +53,8 @@ def _vit_att_fwd( pack_gqa=None, sm_margin=0, sinks=None, + sparse_mask_fine=None, + only_qv=False, ) return o diff --git a/lightllm/models/vit/triton_kernel/flashattention_nopad.py b/lightllm/models/vit/triton_kernel/flashattention_nopad.py index b21464a02b..a9cbcda71d 100644 --- a/lightllm/models/vit/triton_kernel/flashattention_nopad.py +++ b/lightllm/models/vit/triton_kernel/flashattention_nopad.py @@ -203,6 +203,8 @@ def flash_attention_v3_fwd( pack_gqa=None, sm_margin=0, sinks=None, + sparse_mask_fine=None, + only_qv=False, ) return o diff --git a/unit_tests/models/vit/test_flash_attention_forward.py b/unit_tests/models/vit/test_flash_attention_forward.py index 53b70e6836..70849445b4 100644 --- a/unit_tests/models/vit/test_flash_attention_forward.py +++ b/unit_tests/models/vit/test_flash_attention_forward.py @@ -3,6 +3,16 @@ import time import pytest from lightllm.models.vit.triton_kernel.flashattention_nopad import flash_attention_fwd +from lightllm.models.vit.triton_kernel import flashattention_nopad +from lightllm.common.basemodel.attention_vit.fa3.fp import Fa3VitAttBackend + + +_fa3_mark = pytest.mark.skipif( + not flashattention_nopad._flash_attn_v3_available + or not torch.cuda.is_available() + or torch.cuda.get_device_capability()[0] < 9, + reason="FA3 requires sgl_kernel and a Hopper GPU", +) def reference_attention_varlen(q, k, v, cu): @@ -32,7 +42,21 @@ def reference_attention_varlen(q, k, v, cu): @pytest.mark.parametrize("dtype,atol", [(torch.float16, 1e-2), (torch.bfloat16, 2e-2)]) -def test_varlen(dtype, atol, batch=4, heads=8, d=80, device="cuda:0"): +@pytest.mark.parametrize( + "attention_fn", + [ + flash_attention_fwd, + pytest.param( + getattr(flashattention_nopad, "flash_attention_v3_fwd", None), + marks=_fa3_mark, + ), + pytest.param( + Fa3VitAttBackend._vit_att_fwd, + marks=_fa3_mark, + ), + ], +) +def test_varlen(dtype, atol, attention_fn, batch=4, heads=8, d=80, device="cuda:0"): torch.manual_seed(0) lengths = torch.randint(1, 257, (batch,)) max_len = int(lengths.max().item()) @@ -45,10 +69,10 @@ def test_varlen(dtype, atol, batch=4, heads=8, d=80, device="cuda:0"): k = torch.randn_like(q) v = torch.randn_like(q) out_tri = torch.randn_like(q) - flash_attention_fwd(q, k, v, out_tri, cu, max_len) + attention_fn(q, k, v, out_tri, cu, max_len) a = time.time() for _ in range(100): - flash_attention_fwd(q, k, v, out_tri, cu, max_len) + attention_fn(q, k, v, out_tri, cu, max_len) b = time.time() print(f"flash_attention_fwd time: {(b - a) / 100 * 1000:.2f} ms") out_ref = reference_attention_varlen(q, k, v, cu) From 904a295b8cf3fe69a7e47cf3ab9ef19785fb2a3a Mon Sep 17 00:00:00 2001 From: sufubao Date: Sat, 17 Oct 2026 17:17:14 +0800 Subject: [PATCH 2/2] fix: run Triton fallback on the failing FA3 invocation --- .../vit/triton_kernel/flashattention_nopad.py | 4 +- .../vit/test_flash_attention_forward.py | 52 +++++++++++++++++++ 2 files changed, 54 insertions(+), 2 deletions(-) diff --git a/lightllm/models/vit/triton_kernel/flashattention_nopad.py b/lightllm/models/vit/triton_kernel/flashattention_nopad.py index a9cbcda71d..af062c2a5b 100644 --- a/lightllm/models/vit/triton_kernel/flashattention_nopad.py +++ b/lightllm/models/vit/triton_kernel/flashattention_nopad.py @@ -222,8 +222,8 @@ def flash_attention_fwd(q, k, v, o, cu_seqlens, max_seqlen): if _flash_attn_v3_available and (is_hopper() or is_4090()): try: flash_attention_v3_fwd(q, k, v, o, cu_seqlens, max_seqlen) + return except Exception as e: print(f"Failed to use sgl_kernel: {e}") _flash_attn_v3_available = False - else: - _flash_attention_triton_fwd(q, k, v, o, cu_seqlens, max_seqlen) + _flash_attention_triton_fwd(q, k, v, o, cu_seqlens, max_seqlen) diff --git a/unit_tests/models/vit/test_flash_attention_forward.py b/unit_tests/models/vit/test_flash_attention_forward.py index 70849445b4..f736ec4cd6 100644 --- a/unit_tests/models/vit/test_flash_attention_forward.py +++ b/unit_tests/models/vit/test_flash_attention_forward.py @@ -15,6 +15,58 @@ ) +def test_fa3_failure_computes_fallback_on_same_invocation(monkeypatch): + calls = [] + q = torch.zeros(2, 1, 4) + out = torch.full_like(q, -123) + cu = torch.tensor([0, 2], dtype=torch.int32) + + def fail_fa3(q, k, v, o, cu_seqlens, max_seqlen): + calls.append("fa3") + o[0].fill_(99) # A backend may fail after partially writing output. + raise RuntimeError("shape-dependent FA3 failure") + + def fallback(q, k, v, o, cu_seqlens, max_seqlen): + calls.append("triton") + o.fill_(7) + + monkeypatch.setattr(flashattention_nopad, "_flash_attn_v3_available", True) + monkeypatch.setattr(flashattention_nopad, "is_hopper", lambda: True) + monkeypatch.setattr(flashattention_nopad, "flash_attention_v3_fwd", fail_fa3, raising=False) + monkeypatch.setattr(flashattention_nopad, "_flash_attention_triton_fwd", fallback) + flash_attention_fwd(q, q, q, out, cu, 2) + assert calls == ["fa3", "triton"] + assert torch.equal(out, torch.full_like(q, 7)) + assert not flashattention_nopad._flash_attn_v3_available + + # Later calls should use the disabled-backend path without retrying FA3. + out.fill_(-123) + flash_attention_fwd(q, q, q, out, cu, 2) + assert calls == ["fa3", "triton", "triton"] + assert torch.equal(out, torch.full_like(q, 7)) + + +def test_fa3_success_does_not_run_fallback(monkeypatch): + q = torch.zeros(2, 1, 4) + out = torch.full_like(q, -123) + cu = torch.tensor([0, 2], dtype=torch.int32) + + def fa3(q, k, v, o, cu_seqlens, max_seqlen): + o.fill_(9) + + monkeypatch.setattr(flashattention_nopad, "_flash_attn_v3_available", True) + monkeypatch.setattr(flashattention_nopad, "is_hopper", lambda: True) + monkeypatch.setattr(flashattention_nopad, "flash_attention_v3_fwd", fa3, raising=False) + monkeypatch.setattr( + flashattention_nopad, + "_flash_attention_triton_fwd", + lambda *args: pytest.fail("a successful FA3 call must not run fallback"), + ) + flash_attention_fwd(q, q, q, out, cu, 2) + assert torch.equal(out, torch.full_like(q, 9)) + assert flashattention_nopad._flash_attn_v3_available + + def reference_attention_varlen(q, k, v, cu): """ q, k, v : (total_len, n_head, D)