fix(deploy): stop the attention backend report claiming restrictions that do not apply

[rakhimovv]

Follow-up to #9. Two small things in _log_attention_backends that the new guards left slightly out of step — including the one @KraHsu flagged in review as worth a follow-up.

1. The shared-core line can claim a restriction that does not apply

elif shared_backend != "torch":
    shared_backend += " (CUDA fp16/bf16 only; torch_sdpa otherwise)"

ATTENTION_IMPLEMENTATION is not a closed set — initialize_attention_priority() lowercases DIFFSYNTH_ATTENTION_IMPLEMENTATION but does not otherwise check it:

if os.environ.get("DIFFSYNTH_ATTENTION_IMPLEMENTATION") is not None:
    return os.environ.get("DIFFSYNTH_ATTENTION_IMPLEMENTATION").lower()

which is the opposite of how the sibling variable is treated — WAM_ATTENTION_IMPL is validated against _BACKEND_MAP and raises ValueError on an unknown value.

attention_forward's else branch correctly runs any value it does not recognize through torch_sdpa, but the diagnostics then annotate that SDPA run as half-precision-CUDA-only:

$ DIFFSYNTH_ATTENTION_IMPLEMENTATION=sdpa openwam-serve ...
  Wan shared core    : sdpa (CUDA fp16/bf16 only; torch_sdpa otherwise)

That run is plain SDPA on every device and dtype. The empty string is reachable the same way — the lookup guards on is not None, not truthiness — and printed the qualifier attached to no backend name at all. Since #9's purpose was to stop the report claiming a backend the run is not using, this seemed worth closing out.

2. The ActionDiT line carries no qualifier

_flash3 / _flash2 / _sage have always fallen back to SDPA off CUDA, and since #9 they fall back on exactly the same device-and-dtype condition as the Wan paths; _xformers is CUDA-only. The other two lines say so and this one does not, which reads as if the ActionDiT were the odd one out.

What changed

Both ad-hoc conditionals are replaced by one _qualify_backend() whitelist, applied to all three lines. Anything unlisted is reported as-is, which covers plain SDPA and any unrecognized env value.

                      before                                              after
ActionDiT          :  _flash2                                             _flash2 (CUDA fp16/bf16 only; torch_sdpa otherwise)
Wan shared core    :  sdpa (CUDA fp16/bf16 only; torch_sdpa otherwise)    sdpa

The normal case is unchanged: with flash-attn installed and no env override, all three lines read flash_attention_2 (CUDA fp16/bf16 only; torch_sdpa otherwise) / _flash2 (…) as before.

Because the whitelist has to name the ActionDiT closures (_flash2, …) alongside the ATTENTION_IMPLEMENTATION spellings, there is a test pinning those names against components._BACKEND_MAP, so a backend added there later cannot silently drift out of the report.

Tests

Three cases added to tests/test_attention_backend_dispatch.py. Against unpatched server.py they fail 3/3, and the file is 15/15 on this branch.

Verified on Python 3.10.20, torch 2.7.1+cu128, flash_attn 2.8.3.post1, 4×H100 80GB:

full suite (incl. gpu-marked) : 1971 passed, 20 skipped
ruff check / ruff format      : clean
compileall                    : clean

[wayrise]

Both problems are real and I reproduced them exactly as described.

On main, with the env var set to a value attention_forward does not recognize:

$ DIFFSYNTH_ATTENTION_IMPLEMENTATION=sdpa
  Wan shared core    : sdpa (CUDA fp16/bf16 only; torch_sdpa otherwise)
$ DIFFSYNTH_ATTENTION_IMPLEMENTATION=bogus_backend
  Wan shared core    : bogus_backend (CUDA fp16/bf16 only; torch_sdpa otherwise)
$ DIFFSYNTH_ATTENTION_IMPLEMENTATION=
  Wan shared core    :  (CUDA fp16/bf16 only; torch_sdpa otherwise)

and with flash-attn installed and no override, the ActionDiT line was indeed the odd one out:

  ActionDiT          : _flash2
  Video DiT          : flash_attention_2 (CUDA fp16/bf16 only; torch_sdpa otherwise)
  Wan shared core    : flash_attention_2 (CUDA fp16/bf16 only; torch_sdpa otherwise)

All four cases read correctly on this branch. The WAM_ATTENTION_IMPL / DIFFSYNTH_ATTENTION_IMPLEMENTATION asymmetry you point to is real too — the former raises ValueError against _BACKEND_MAP, the latter takes any string.

Claim	Result
15/15 on this branch	Confirmed
3/3 fail against unpatched server.py	Confirmed
Full suite	Confirmed — 1971 passed, 20 skipped
ruff check / ruff format	Clean

Collapsing three ad-hoc conditionals into one whitelist is the right shape, and test_action_dit_backend_names_match_components is a good catch — pinning the naming convention rather than the closure objects is the only thing that works when CI has none of the libraries installed.

One nit, non-blocking

The whitelist merges two namespaces into one flat set, so a name from one can match an entry meant for the other. Since ATTENTION_IMPLEMENTATION is an open set — which is the premise of this PR — the ActionDiT closure names are reachable through it:

$ DIFFSYNTH_ATTENTION_IMPLEMENTATION=_flash2
  Wan shared core    : _flash2 (CUDA fp16/bf16 only; torch_sdpa otherwise)

>>> resolve_implementation(q)          # '_flash2'
>>> attention_forward(q, q, q)         # runs fine on cpu/fp32 — it is plain SDPA

That is the same defect this PR fixes, just entering through the other namespace's names. Nobody sets the env var to an ActionDiT internal, so it is not worth blocking on — but if you want it airtight, giving _qualify_backend the namespace (two small sets consulted separately, or a namespace argument) closes it without changing anything else.

Related, and purely cosmetic: DIFFSYNTH_ATTENTION_IMPLEMENTATION= still logs Wan shared core : with a blank name. Strictly an improvement over annotating that blank with a restriction, and arguably the honest "report as-is" answer, but something like torch_sdpa (unrecognized: "") would read better in a log. Entirely your call.

Neither of these needs to hold up the merge. Approving.

[rakhimovv]

Thanks for the thorough reproduction, and the nit is a good catch — I've taken it rather than leaving it, because it turned out to be the same defect class the PR is about and the fix also closed a gap I'd found on my own side.

Pushed as c5b6b81.

Namespaces split. _qualify_backend now takes the calling subsystem's own table (_WAN_BACKENDS or _ACTION_BACKENDS) instead of consulting one flat set, so a name is annotated only where it means something. Your exact case:

$ DIFFSYNTH_ATTENTION_IMPLEMENTATION=_flash2
  Wan shared core    : _flash2          # was: _flash2 (CUDA fp16/bf16 only; ...)

The table values are now the restriction note itself rather than set membership, which also lets the drift tests assert which restriction a backend carries — a backend classified into the wrong one used to pass.

And the counterpart pin you didn't ask for. Writing the split made obvious something I should have caught earlier: test_action_dit_backend_names_match_components pinned only one of the three sources that generate backend names. fused_backend_name() and initialize_attention_priority() were both unpinned, so a backend added to either would silently lose its qualifier — the same failure mode, one level up. test_wan_backend_names_match_their_sources closes that.

Both new tests are mutation-tested: flattening the two tables back into one fails test_a_name_from_one_namespace_is_not_qualified_in_the_other, and dropping sage_attention from _WAN_BACKENDS fails test_wan_backend_names_match_their_sources. Neither passes vacuously.

On the blank-name cosmetic — I left it as-is. DIFFSYNTH_ATTENTION_IMPLEMENTATION= still logs Wan shared core : with an empty name. Your torch_sdpa (unrecognized: "") reads better, but it means _qualify_backend inventing a name for a value it was handed, and the honest "report what is set" answer felt like the right shape for a function whose whole purpose is to stop the report claiming things. Happy to add it if you'd rather — it's a two-line change and I have no strong attachment.

Verified on Python 3.10.20, torch 2.7.1+cu128, flash_attn 2.8.3.post1, 4×H100: full suite 1973 passed, 20 skipped, ruff check / ruff format --check clean. The normal path is unchanged — with flash-attn installed and no override, all three lines still read flash_attention_2 (CUDA fp16/bf16 only; torch_sdpa otherwise) / _flash2 (…) as before.

Sorry to move it under an approval you'd already given — re-review whenever suits, or say the word and I'll drop the second commit if you'd rather land the first as approved.

[wayrise]

Re-reviewed c5b6b81. No need to apologise for the dismissed approval — taking the nit was the right call over leaving it.

The namespace split does what it says. My case, on this head:

$ DIFFSYNTH_ATTENTION_IMPLEMENTATION=_flash2
  Wan shared core    : _flash2
$ DIFFSYNTH_ATTENTION_IMPLEMENTATION=_xformers
  Wan shared core    : _xformers

Both plain now, and the ordinary path is unchanged — with flash-attn installed and no override, all three lines still carry their qualifier.

Making the table values the restriction note rather than set membership is a real improvement on its own: classifying a backend into the wrong bucket used to pass, and now it cannot.

I re-ran your two mutations rather than take them on trust, and both hold:

mutation	result
drop sage_attention from _WAN_BACKENDS	1 failed
flatten the two tables back into one	2 failed
unmutated control	17 passed

check	result
tests/test_attention_backend_dispatch.py	17 passed
full suite	1973 passed, 20 skipped
ruff check / ruff format --check	clean

On the blank name — agreed, leave it. A function whose job is to stop the report claiming things is the wrong place to invent a name for a value it was handed, and "" is the honest answer to "what is set". If it ever bothers anyone in a log it can be dealt with where the env var is read, which is the layer that actually knows the value was junk.

Approving.

[rakhimovv]

Thanks for the re-review, and for re-running the mutations rather than taking them on trust.

Closing the one thing I left open: I said I'd add your torch_sdpa (unrecognized: "") if you preferred it, and on reflection that was me handing back a decision rather than making one. I'm not adding it, and I've filed #12 instead.

The blank name is a symptom of DIFFSYNTH_ATTENTION_IMPLEMENTATION being taken unvalidated, and while writing that up I found it also produces a genuine crash, which I hadn't tested before:

$ DIFFSYNTH_ATTENTION_IMPLEMENTATION=flash_attention_3     # FA3 not installed
resolve_implementation   -> flash_attention_3
attention_forward: NameError: name 'flash_attn_interface' is not defined

WAM_ATTENTION_IMPL=flash3 warns and falls back on the same box. So the empty-name log line, the silently-accepted typo, and that NameError are all one missing validation, and dressing the first one up in the logger would have made the report look complete while the crash stayed. #12 has all three modes and a proposed fix; the one semantics question in it is genuinely yours, so I've asked rather than guessed.

Nothing further from me on this PR — it's ready whenever you are.