[Transformers backend] Find attention with a fuser and attach vLLM's layer to it - #54941
Conversation
Signed-off-by: Thomas Ortner <boh@zurich.ibm.com>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
…ntation` Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
|
Documentation preview: https://vllm--54941.org.readthedocs.build/en/54941/ |
mergify
Bot
added
documentation
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
|
/ci run |
|
✅ Triggered Buildkite CI #86894 for commit |
hmellor
changed the title
hmellor
changed the title
|
This pull request has merge conflicts that must be resolved before it can be |
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
|
/ci run |
|
✅ Triggered Buildkite CI #86905 for commit |
hmellor
left a comment
hmellor
left a comment
There was a problem hiding this comment.
If the CI passes this should be good
|
✅ @bohnstingl, CI is now available for this PR.
|
Signed-off-by: Thomas Ortner <boh@zurich.ibm.com>
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Team Run ID: 📒 Files selected for processing (14)
🚧 Files skipped from review as they are similar to previous changes (14)
Included review availability: Your plan provides up to 10 included reviews per hour; 9 remain after this review. 📝 SummarySummary by CodeRabbit
WalkthroughThe Transformers backend now detects attention interface calls, supports multiple fusers per module, tracks dispatchers by layer, derives attention scales, and attaches vLLM attention instances to serving modules. Tests and custom-model documentation cover the updated requirements. ChangesTransformers attention fusers
Estimated code review effort: 4 (Complex) | ~45 minutes Merge Risk: ⚪ Minimal · up to This change adds Transformers attention fusing and scale resolution with targeted test coverage reported. No concrete merge-blocking risk remains in the supplied change context. Sequence Diagram(s)sequenceDiagram
participant TransformersBackend
participant AttentionFuser
participant AttentionModule
participant vllm_attention_forward
participant VLLMAttention
TransformersBackend->>AttentionFuser: register attention fuser by layer
TransformersBackend->>AttentionModule: attach VLLMAttention
AttentionModule->>vllm_attention_forward: dispatch attention interface
vllm_attention_forward->>AttentionModule: resolve attached attention
vllm_attention_forward->>VLLMAttention: execute attention with module inputs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Full details: Docstring CoverageExplanation Docstring coverage is 60.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 65 functions across 13 files. (1 skipped: 1 unsupported.) Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Signed-off-by: Thomas Ortner <boh@zurich.ibm.com>
Signed-off-by: Thomas Ortner <boh@zurich.ibm.com>
|
/ci run |
|
✅ Triggered Buildkite CI #87065 for commit |
Signed-off-by: Thomas Ortner <boh@zurich.ibm.com>
…ttn-module Signed-off-by: Thomas Ortner <boh@zurich.ibm.com>
This reverts commit 4cc0f6e. Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
|
I've added the recursion back because it's necessary |
|
/ci run |
|
✅ Triggered Buildkite CI #87069 for commit |
|
/ci run |
|
✅ Triggered Buildkite CI #87075 for commit |
|
/ci retry |
|
✅ Queued 1 failed job(s) for retry in Buildkite CI #87075. |
|
/ci run |
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
|
✅ Triggered Buildkite CI #87250 for commit |
|
/ci run |
|
✅ Triggered Buildkite CI #87258 for commit |
github-project-automation
Bot
moved this from In review
to Shipped
in Transformers modeling backend
…layer to it (vllm-project#54941) Signed-off-by: Thomas Ortner <boh@zurich.ibm.com> Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Signed-off-by: Jyotirmoy Roy <jyotirmoyroy649@gmail.com>
2 participants
Purpose
Teach the Transformers modeling backend to find each layer's attention with a fuser, attach
vLLM's attention layer to the HF module that dispatches it, and resolve the softmax scale at
construction.
create_attention_instancesbuilt oneAttentionper layer intoself.attention_instances, a plaindict.nn.Module.__setattr__does not register a dict, sonone of those layers was part of the module tree, and
vllm_attention_forwardreached its ownwith
attention_instances[module.layer_idx]. The MLA path already worked around half of this byattaching its instance as
mla_module._vllm_mla_attn, precisely so the layer appeared innamed_modules()and ran itsprocess_weights_after_loading.AttentionFuser(new). A module dispatches attention if itsforwardbindsALL_ATTENTION_FUNCTIONS.get_interface(...)to a local and calls it exactly once. That is readfrom the forward's source, not from the
fxtrace:traceis deliberately partial-tolerant andreturns whatever graph it managed, so a missing interface node does not mean the module has no
dispatch. Keying off
self.fuserswould be wrong for the same reason a QKV fuser is not theright handle -- a model can dispatch through the interface without its projections fusing at all.
The fuser also carries the
scaling=expression it found, and answerslayer_indexfrom themodule's own
layer_idx.Fusers are now plural per module.
get_fusersreturns every fuser that applies: at most onethat redefines the forward (
redefines_forward, mutually exclusive because each rewrite startsfrom the original source) plus any number that leave it alone.
AttentionFuseris the first ofthe latter kind, so a Llama attention comes back as
[QKVFuser, AttentionFuser]and a Gemma 4attention -- whose QKV fusion is rejected -- as
[AttentionFuser]. This also drops a special case:an
MLAFuserthat failsvalidateused to be rebuilt as a bare attention marker, and now simplyfalls out of the list while the
AttentionFuserstands on its own.This fixes two latent bugs and removes one requirement:
impl.scalewas written per forward, which one compiled artifact per layer class cannotpreserve. Every layer was built with the Llama default
head_size**-0.5and corrected fromthe HF module's
scalingkwarg on every call. That correction only survives becausemodule.layer_idxis anintbaked into the graph, so Dynamo compiles the stack layer bylayer. The moment one artifact is reused across layers -- the point of
_USE_LAYERNAME/LayerNamehoisting -- only the first layer's Python frame runs and layers 1..N silently keepthe default. Wrong for any model whose scale is not derived from
head_size. The scale is nowread from the
scaling=argument the module hands the interface, once, at construction.Reading the argument rather than probing
module.scalingalso matters on its own: OPT applieshead_size**-0.5to the query itself and then declaresscaling=1.0, so probing the attributewould scale twice.
getattr(module, "attn")is the same expression in every layer, so one graph can serve thewhole stack.
attention_instanceskwarg no longer has to be drilled down. With nothing to pass, theregistry gate relaxes from
is_backend_compatible()(_supports_attention_backend, documentedupstream as "fully pass the kwargs through all modules up to the Attention layer") to
_can_set_attn_implementation(), which asks only that the model dispatches through theinterface.
Attaching also makes the backend's module tree the same shape as an in-tree model's, so anything
that walks it --
process_weights_after_loading, a state dict,model.to(device)-- findsattention where it expects to. The attribute is
attn, the name an in-tree model uses(
LlamaAttention.attn), which is also whatmaybe_remap_kv_scale_namewrites(
.self_attn.attn.{k,v}_scale), so a checkpoint's FP8 KV-cache scales land on it with no mapperentry.
_vllm_mla_attnis gone; it was never read outsidecreate_attention_instances.Test Plan
New tests in
tests/models/transformers/test_backend.py, alongside the device-free ones alreadythere. They need no accelerator, no weights and no network: the HF stacks are built from
AutoConfig.for_model(...)on themetadevice.pytest -q tests/models/transformers/fusers/ pytest -q tests/models/transformers/test_backend.py \ -k "attention_dispatch or attention_scale or layer_index"test_attention_dispatch_is_matched-- exactly the decoder layers' attention modules match anAttentionFuser, overllama,gemma4_textanddeepseek_v3(MLA).test_attention_layer_index_is_the_modules_own-- LongCat Flash gives each decoder layer twoattention sublayers numbered
2iand2i + 1, sonum_hidden_layersis twice the length of thestack and the enclosing layer's position is not the index the KV cache is keyed by. Also pins
that
validateis what excludes a module whose config was not patched to dispatch to vLLM, whichis how a vision tower is left to Transformers.
test_attention_scale_is_the_declared_one-- the scale is the module's, not the default.gemma4_text(learnable per-dim query weight, so it declares 1.0) anddeepseek_v3(yarnmscale) are what make this more than a tautology.test_attention_scale_is_the_argument_not_the_attribute-- OPT, where the argument (1.0) and theattribute (
head_size**-0.5) disagree.The existing fuser suites cover the registry change;
test_linear.py's QKV and packed-QKV numericstests now drive the real
vllm_attention_forwardthrough the attached.attn.Test Result
The file's other device-free tests -- embedding replacement and multimodal component marking --
pass alongside them, unchanged (3 passed).
Not yet run, and needed before merge. Everything requiring an accelerator: the rest of
tests/models/transformers/,tests/models/test_initialization.py -k Transformers, and modelevals. This changes attention construction for every model on the backend, so
tests/evals/orvllm benchresults belong here before it is reviewed for correctness. Tensor parallel is alsounexercised.
Notes for reviewers
Attentioninstances visible tonamed_modules()by registering them in annn.ModuleListrather than attaching them to themodule that dispatches. The two need reconciling; attaching is the one that also removes the
dict index from the traced graph and gives kv-scale remapping the name it expects.
fx_utilsalready records the interface call as a leafnode, which would be the obvious hook, but
tracereturns a partial graph on failure and a modelcan stop tracing before the dispatch. Reading the forward's source is also what Transformers
itself does in
_can_set_attn_implementation, and it is where thescaling=argument has to comefrom regardless.
without its projections fusing at all, in which case it would never have been attached.
the dict, as does a layer whose modules claim the same
layer_idx. The latter was already broken:two modules sharing an index would have shared one
Attention.Duplicate check
gh pr list --state open --search "transformers backend attention in:title,body"returns no otherPR changing how the backend builds or reaches its attention layers. #52156 is the only related one
and is called out above.
Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.