From 2b33c1527782e77fea9f399b687c9e828cacb0a6 Mon Sep 17 00:00:00 2001 From: Pedro Augusto <16762743+pedronaugusto@users.noreply.github.com> Date: Wed, 22 Apr 2026 00:14:53 +0100 Subject: [PATCH] sparse/attention: drop FUSED_ATTN_MAX_SEQLEN threshold to match CUDA-path precedent MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The 256-cap inside the flex_gemm_sparse_attn branch was a hold-over from the early naive Metal kernel (round 3). The current backend is flash- attention-v2 with simdgroup matmul + simd-shuffle softmax row reductions and wins at every measured shape including max_seqlen=2048. The CUDA backends above (xformers / flash_attn / flash_attn_3) never fork on max_seqlen — match that precedent. Safety-valve preserved as FLEX_GEMM_ATTN_MAX_SEQLEN=N env var: when set, falls back to SDPA-padded above the cap. Useful only on PyTorch builds where the Accelerate-SDPA-CPU-bounce happens to win at a specific shape (measured crossover sits beyond 768 on fp32, higher on fp16). --- .../modules/sparse/attention/full_attn.py | 32 ++++++++++++------- 1 file changed, 21 insertions(+), 11 deletions(-) diff --git a/trellis2/modules/sparse/attention/full_attn.py b/trellis2/modules/sparse/attention/full_attn.py index dceceeb..eceb6c0 100644 --- a/trellis2/modules/sparse/attention/full_attn.py +++ b/trellis2/modules/sparse/attention/full_attn.py @@ -213,19 +213,29 @@ def sparse_scaled_dot_product_attention(*args, **kwargs): out = flash_attn_3.flash_attn_varlen_func(q, k, v, cu_seqlens_q, cu_seqlens_kv, max_q_seqlen, max_kv_seqlen) elif config.ATTN == 'flex_gemm_sparse_attn': # Fused variable-length sparse attention (mtlgemm Metal kernel). - # Keeps everything on MPS — avoids the MPS->CPU->MPS round-trip that - # SDPA-padded forces on PyTorch builds where new_zeros for MPS - # fp16/fp32 is broken. Wins measurably when max_seqlen is small-ish - # (trellis2 decoder's typical case); at larger max_seqlen the current - # naive per-thread-serial-KV kernel is slower than Accelerate-backed - # SDPA, so we fall back. Threshold is conservative; measured - # crossover was around 256-512 on M3 Max. - FUSED_ATTN_MAX_SEQLEN = 256 + # Dispatch is unconditional, matching the CUDA backends above + # (xformers / flash_attn / flash_attn_3) which never fork on + # max_seqlen. The underlying kernel is flash-attention-v2 with + # simdgroup_matrix_multiply_accumulate for Q@K^T and P@V, and + # parallelized online-softmax row reductions via simd_shuffle_xor. + # + # Optional safety-valve: FLEX_GEMM_ATTN_MAX_SEQLEN=N falls back to + # the SDPA-padded path when max(q_seqlen, kv_seqlen) > N. Useful on + # PyTorch builds where the Accelerate-SDPA-CPU-bounce is actually + # faster at the specific shapes the user hits (rare; measured + # crossover sits beyond 768 on fp32 and likely higher on fp16). + import os as _os + _cap_env = _os.environ.get('FLEX_GEMM_ATTN_MAX_SEQLEN', '0') + try: + _cap = int(_cap_env) + except ValueError: + _cap = 0 if num_all_args == 1: q, k, v = qkv.unbind(dim=1) elif num_all_args == 2: k, v = kv.unbind(dim=1) - if max(max(q_seqlen), max(kv_seqlen)) <= FUSED_ATTN_MAX_SEQLEN: + _use_fused = (_cap <= 0) or (max(max(q_seqlen), max(kv_seqlen)) <= _cap) + if _use_fused: import flex_gemm, math scale = 1.0 / math.sqrt(q.shape[-1]) csq = torch.cat([torch.tensor([0]), torch.cumsum(torch.tensor(q_seqlen), 0)]).int().to(device) @@ -235,8 +245,8 @@ def sparse_scaled_dot_product_attention(*args, **kwargs): csq, cskv, max(q_seqlen), max(kv_seqlen), scale, ) else: - # Fall through to the SDPA implementation below by mirroring its - # preamble here (we already have q/k/v unbound). + # FLEX_GEMM_ATTN_MAX_SEQLEN safety-valve active — fall back to + # SDPA-padded. Mirror the SDPA preamble (q/k/v already unbound). import torch.nn.functional as F_attn N_b = len(q_seqlen) max_q = max(q_seqlen); max_kv = max(kv_seqlen)