Skip to content

[Feature] Declare the 'fa' backend's sm_90 requirement in BackendInfo so an explicit --attention-backend fa fails at config time, not in graph capture #442

Description

@salekseev

Applies to

attention/fa.py and the BackendInfo capability matrix in attention/__init__.py. Verified on main @ fb7f732.

What problem does this solve

An explicit --attention-backend fa on sm_89 passes config validation, loads the whole model, and then dies inside CUDA-graph capture with a raw CUDA error:

CUDA error (_deps/repo-flash-attention-src/hopper/flash_fwd_launch_template.h:203): invalid argument
Backend supervisor: backend worker freetoken-TP0-scheduler exited during load

That is roughly 40 seconds of expert loading before anything complains, and the message does not point at the flag that caused it.

The information needed to refuse it early is already in the tree, just not in the capability matrix:

  • _resolve_auto_attention_backend gates the fa,fi tier on is_sm90_family(), so auto never picks fa below Hopper.
  • fa's BackendInfo (attention/__init__.py:73) declares only requires_sgl_kernel=True. There is no arch field, where trtllm has requires_sm100=True.
  • FlashAttentionBackend.__init__ (fa.py:38) rejects only compute capability 12.x, and fa.py:52 then sets self.version = 4 if is_sm100_supported() else 3, so sm_89 asks for FA3.

The underlying reason FA3 cannot run there is on sglang's side: the published sglang-kernel wheels contain no sm_89 cubin for flash_ops (cuobjdump --list-elf gives sm_80 sm_86 sm_90 for both 0.4.5 and 0.4.6.post1), even though their is_fa3_supported() accepts major 8. I have filed that separately with them. This issue is only about failing early and legibly on our side.

Proposed solution

Give BackendInfo an sm_90 requirement and set it on fa, matching what the wheels actually contain:

requires_sm90: bool = False

_backend_requirements_met already has the shape for this -- it checks requires_sm100 against is_sm100_family() -- so the change is one field, one check, and one flag on the fa entry. _validate_attention_backend_choice would then refuse an explicit --attention-backend fa on Ada at config time with the existing message style, and auto behaviour is unchanged because the tier is already gated.

Alternatives considered

  • Leave it to sglang. Their is_fa3_supported() is the root cause and I have reported it. But even after a fix there, the capability matrix is the right place for FreeToken to express "this backend needs Hopper", and the auto tree already encodes exactly that belief.
  • Raise in FlashAttentionBackend.__init__ alongside the existing 12.x check. That works but fires later than config validation and duplicates knowledge the matrix should hold.
  • Do nothing. Defensible, since auto protects everyone who does not pass the flag explicitly. The cost is that anyone who does pass it burns a full model load to find out.

For reference on what the alternatives are worth on this hardware: fi measured 117.23 tok/s decode at 1200 tokens and 105.04 at 32k, against 115.34 and 91.24 for triton at the same settings with bf16 KV. So fi is the right default on sm_89 and fa is not leaving anything on the table here, which is partly why this is a papercut rather than a blocker.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions