[Fusion] Manual ActivationQuantFusionPass initial application - #51415
Conversation
Start the ActivationQuantFusionPass manual-fusion migration (RFC vllm-project#43224) on the producer side of the QuantizedActivation contract (vllm-project#44260). Adds maybe_fused_act_quant: given an activation and the linear it feeds, it emits a QuantizedActivation via the fused silu_and_mul_quant kernel when the linear advertises a consumable input_quant_key, and falls back to the plain activation otherwise. Llama's MLP forward calls it for down_proj; the change to model code is one line. When manual fusion fires the silu_and_mul pattern is already consumed, so the compiler ActivationQuantFusionPass finds nothing and the two never double-fuse. Only (SiluAndMul, kFp8StaticTensorSym) is registered for now. Smoke tested on RedHatAI/Llama-3.2-1B-Instruct-FP8 (default compiled + CUDA-graph path): all down_proj layers take the fused branch and generations are coherent. AI assistance (Claude) was used; all changed lines are human-reviewed. Signed-off-by: mgoin <mgoin64@gmail.com> Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| ) | ||
|
|
||
|
|
||
| # (activation module type, consumer input_quant_key) -> fused producer. |
There was a problem hiding this comment.
do you think we should build in some oracle for this selection up front where the actual implementation of the quantized activation function can be swapped?
There was a problem hiding this comment.
Yeah it is reasonable to suspect for future optimizations
| def forward(self, x): | ||
| x, _ = self.gate_up_proj(x) | ||
| x = self.act_fn(x) | ||
| x = maybe_fused_act_quant(self.act_fn, x, self.down_proj) |
There was a problem hiding this comment.
I think that instead of passing self.down_proj: LinearBase, it might make more sense to expose a method from LinearBase that can be called here to extract what is needed
I think that will help make the interface clearer
There was a problem hiding this comment.
Depending on the quantization used, we need to extract different things from the Linear layer. IMO it is more clear to my eye that this function might have an affect on self.down_proj within the context of the forward function, so I prefer just passing it in
|
Hi @mgoin, I’d like to help validate this manual activation+quant fusion work. I’m planning to run TP=1/TP=2 correctness checks and A/B tests comparing the base commit and this PR with compiler activation-quant fusion enabled and disabled. Please let me know if there’s a preferred test setup or specific coverage you’d find most useful. |
Signed-off-by: ElizaWszola <ewszola@redhat.com>
Signed-off-by: ElizaWszola <ewszola@redhat.com>
Signed-off-by: ElizaWszola <ewszola@redhat.com>
Signed-off-by: ElizaWszola <ewszola@redhat.com>
|
/ci run |
|
✅ Triggered Buildkite CI #83767 for commit |
|
/ci run |
1 similar comment
|
/ci run |
|
✅ Triggered Buildkite CI #84616 for commit |
Signed-off-by: ElizaWszola <ewszola@redhat.com>
|
/ci run |
|
✅ Triggered Buildkite CI #86403 for commit |
|
@mgoin I bisected a CI IMA to this change: https://buildkite.com/vllm/ci/builds/87269/list?jid=01a06dce-0d6b-4b62-8db0-9853df2b22d9&tab=output |
5 participants
Starts the
ActivationQuantFusionPassmanual-fusion migration (RFC #43224, specific tracker #43501) on the producer side of theQuantizedActivationcontract (#44260).Adds
maybe_fused_act_quant: given an activation and the linear it feeds, it emits aQuantizedActivationvia the fusedsilu_and_mul_quantkernel when the linear advertises a consumableinput_quant_key, and falls back to the plain activation otherwise. Llama's MLP forward calls it fordown_proj— the model-code change is one line. When manual fusion fires thesilu_and_mulpattern is already consumed, so the compilerActivationQuantFusionPassfinds nothing and the two never double-fuse. Only(SiluAndMul, kFp8StaticTensorSym)is registered for now; other schemes/activations are one table row each.Not a duplicate: builds on the landed contract (#44260) and is the activation-quant counterpart to the AR+RMSNorm manual fusion in #45855; no open PR covers the producer side.
Test
Smoke tested on
RedHatAI/Llama-3.2-1B-Instruct-FP8(default compiled + CUDA-graph path): all 16mlp.down_projlayers take the fused branch (input_quant_key == kFp8StaticTensorSym) and generations are coherent — theQuantizedActivationflows through the compileddown_proj.forwardintoapply_weightswith no graph break. TP2 coherence and an A/B vs. the compiler pass are follow-ups.AI assistance (Claude) was used; all changed lines are human-reviewed.