Conversation
Introduce an opt-in path for model code to pre-quantize activations for quantized linears, so RMSNorm + input-quant fusion can be expressed directly in model code rather than relying on torch.compile passes. - Add QuantizedActivation + rms_norm_input_quant helper. - Split CutlassNvFp4LinearKernel.apply_weights into apply_weights (chain) and apply_quantized (pre-quantized entry point). - Have CompressedTensorsW4A4Fp4 advertise layer.input_quant_key and implement quantize_input; route QuantizedActivation through to the kernel's apply_quantized. - Wire Qwen3 decoder layer to call rms_norm_input_quant before qkv_proj and gate_up_proj. Only the Cutlass NVFP4 kernel opts in; FlashInfer/fbgemm fall back to the legacy in-apply quant path. Other quant methods and other models are unchanged. Requires --enforce-eager (or piecewise compile) to take effect; the isinstance(QuantizedActivation) dispatch breaks the torch.compile graph by design, which is the direction this rework is moving toward. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: mgoin <mgoin64@gmail.com>
…env var Extend the manual-input-quant hoisting prototype to compressed_tensors static-per-tensor FP8 (kFp8StaticTensorSym) and wire the Llama decoder layer to use it. Smoke-tested on RedHatAI/Llama-3.2-1B-FP8 with byte-identical output vs. the legacy path. - Gate scheme opt-in on VLLM_HOIST_INPUT_QUANT=1; default off keeps the torch.compile + rms_quant_fusion pass active as today. - CompressedTensorsW8A8Fp8: opt in only when static-tensor activation; quantize_input calls scaled_fp8_quant(layer.input_scale), apply_weights routes QuantizedActivation by passing the FP8 tensor through (kernel already has a "skip quant if x.dtype == fp8" branch). - CompressedTensorsW4A4Fp4: same gate retroactively applied so the NVFP4 path also defaults off. - llama.py LlamaDecoderLayer.forward: call rms_norm_input_quant before qkv_proj and gate_up_proj. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: mgoin <mgoin64@gmail.com>
Add a scheme-level rms_norm_quantize_input that the helper prefers over the chained norm+quantize path. CompressedTensorsW8A8Fp8 implements it by calling torch.ops._C.rms_norm_static_fp8_quant / fused_add_rms_norm_static_fp8_quant directly, so manual-fusion mode hits the same C++ kernel as the compile-fusion path. CompressedTensorsLinearMethod gains a method-level forwarder that delegates to the scheme's fused entry point if present, otherwise falls back to chained norm + scheme.quantize_input. NVFP4 keeps the chained path since no real fused kernel exists yet. Tested on RedHatAI/Llama-3.2-1B-FP8 across all four mode combinations (eager/compile x HOIST=0/1) with byte-identical generated text; verified under fullgraph torch.compile that Dynamo traces into the new path. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: mgoin <mgoin64@gmail.com>
Both manual-fusion and compile paths produce byte-identical output and both reach the same fused C++ kernel, so the env-var gate is just dead weight. Each scheme keeps its own narrow correctness gate (NVFP4: only when the kernel exposes apply_quantized; FP8: only static-per-tensor), so non-supported configurations still fall through to the legacy path. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: mgoin <mgoin64@gmail.com>
…o kernel Strip the NVFP4 + Qwen3 changes to focus the prototype on static FP8. Move the QuantizedActivation-vs-Tensor dispatch from CompressedTensorsW8A8Fp8.apply_weights down into FP8ScaledMMLinearKernel.apply_weights so the scheme and the CompressedTensorsLinearMethod are pure passthroughs; the isinstance check lives in exactly one place. Drop the now-unused CompressedTensorsW8A8Fp8.quantize_input and CompressedTensorsLinearMethod.quantize_input forwarder, and collapse rms_norm_quantize_input on the method to a one-line delegation. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: mgoin <mgoin64@gmail.com>
Make quant_fusion agnostic of the LinearMethod/Scheme abstractions: it
reads layer.input_quant_key and dispatches directly to the right fused
op (today: kFp8StaticTensorSym → torch.ops._C.{rms_norm,fused_add_rms_norm}_static_fp8_quant).
Remove the now-unused rms_norm_quantize_input methods from
CompressedTensorsW8A8Fp8 and CompressedTensorsLinearMethod.
Co-authored-by: Claude <noreply@anthropic.com>
Signed-off-by: mgoin <mgoin64@gmail.com>
Make the rms_norm_input_quant signature kwarg-explicit so model code documents what's being pulled from the downstream linear instead of handing the whole module to the helper to introspect. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: mgoin <mgoin64@gmail.com>
…all site" This reverts commit 5c84221.
Make rms_norm_input_quant transparent to AllReduce: when residual is not None and tp_size>1 the helper calls tensor_model_parallel_all_reduce on the input before the (potentially fused) add+rms_norm+quant op. The discriminator residual-is-None matches the decoder-layer convention: layer 0's input_layernorm has no residual and no AR; every other site has both. Llama and Qwen3 decoder layers now construct o_proj and the MLP's down_proj with reduce_results=False so the AR happens inside the helper. At tp=1 this is a strict no-op (RowParallelLinear already skipped AR when tp_size=1). Smoke-tested RedHatAI/Llama-3.2-1B-FP8 on tp=1 (output unchanged) and tp=2 (correct, coherent output in both eager and torch.compile). Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: mgoin <mgoin64@gmail.com>
Drop the residual-is-not-None heuristic in favor of inspecting the producer linear directly: rms_norm_input_quant now AR's iff a prev_linear is supplied and that linear has tp_size>1 and reduce_results=False. disable_tp=True (which leaves a linear at tp_size=1) short-circuits cleanly. Route LlamaModel/Qwen2Model's final norm through the same helper, passing the last local layer's down_proj as prev_linear. Plain Qwen2 keeps reduce_results=True so the helper is a no-op for it; Llama and Qwen3 set reduce_results=False so the helper completes the AR before normalizing. This fixes a correctness bug where the final norm was operating on a per-rank partial sum, producing wrong logits at tp>1. Eager and compile outputs now agree on Llama-3.2-1B-FP8 at tp=2. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: mgoin <mgoin64@gmail.com>
When the manual fusion path needs an AllReduce, prefer the FlashInfer single-kernel op (torch.ops.vllm.flashinfer_trtllm_fused_allreduce_norm) over the two-kernel chain. Falls back to the chain when FlashInfer isn't loaded or when the input isn't 2D. The op needs a max_token_num for its workspace; capture it on the layer in process_weights_after_loading (where vllm_config is still in context) rather than calling get_current_vllm_config() at forward time. Tested on RedHatAI/Llama-3.2-1B-FP8 tp=2 eager + compile: identical generated text in both modes (workspace log line "Initialized FlashInfer Allreduce norm quantization fusion workspace with backend=trtllm" confirms the fused path fires). Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: mgoin <mgoin64@gmail.com>
quant_fusion now owns a `QUANT_IMPLS` dispatch table (FP8 static per-tensor, NVFP4) and a `FusedAllReduceRMSQuant` module that pre- resolves the impl + max_token_num at construction time. Decoder layers instantiate one per (norm, consumer-linear) pair instead of calling a helper at forward time — keeps the fusion choice on the fused op, not smeared across the linear. FlashInfer wired directly (CUDA-only top-level import; no try/except fallback chain). Linear methods set `layer.input_quant_key` once in `create_weights` — both compressed-tensors W4A4 NVFP4 and ModelOpt NVFP4 — so the fused-op constructor in the decoder layer can resolve its dispatch before forward. NVFP4 path: `_nvfp4` allocates the FP4 quant_out + swizzled scale_out via `create_fp4_output_tensors`, calls `kARResidualRMSNormFP4Quant` with `scale_factor=input_global_scale_inv`, packages into a `QuantizedActivation`. `FlashInferCutlassNvFp4LinearKernel` now unpacks `QuantizedActivation` on the fast path and skips the redundant `scaled_fp4_quant` call. Smoke-tested on B300 (sm_103) with Llama-3.2-1B FP8 and inference-optimization/Llama-3.2-1B-Instruct-NVFP4 at TP=1/2 in eager + compile; FlashInfer fused workspace initializes and the fused single-kernel fires for both FP8 and NVFP4 paths. Signed-off-by: mgoin <mgoin64@gmail.com>
Signed-off-by: mgoin <mgoin64@gmail.com>
Signed-off-by: mgoin <mgoin64@gmail.com>
Manual fusion (fused_ar_rms_norm_quant) keys off layer.input_quant_key to decide whether the upstream RMSNorm should emit the pre-quantized activation. That key was hard-coded in the W8A8 FP8 scheme and only set on the scheme object (not the layer) for the NVFP4 schemes. Both have problems: - On non-CUDA, the W8A8 scheme would still mark the layer with kFp8StaticTensorSym and dispatch into the CUDA-only flashinfer AR+RMS+FP8 path. - The NVFP4 schemes set self.input_quant_key on the scheme rather than the layer, so getattr(consumer, "input_quant_key", None) at the fusion call site always returned None and the fast path was dead. Move the source of truth to the kernel: - FP8ScaledMMLinearKernel.input_quant_key() defaults to None. CutlassFP8 and FlashInferFP8 override to return kFp8StaticTensorSym when the config is static per-tensor. All other kernels (AITER, ROCm, CPU, XPU, Marlin, PerTensorTorch) inherit None. - Fp8BlockScaledMMLinearKernel.input_quant_key() returns None (block scales are dynamic per-group; nothing to hoist). - NvFp4LinearKernel.input_quant_key() was already in place; FlashInferCutlassNvFp4LinearKernel returns kNvfp4Dynamic. - Schemes assign layer.input_quant_key = kernel.input_quant_key() in create_weights when the result is non-None. Result: non-CUDA platforms never set layer.input_quant_key, so the manual-fusion helper falls through to AR + norm.forward. Signed-off-by: mgoin <mgoin64@gmail.com>
The previous import block assumed FlashInfer was always available on CUDA. ``allreduce_rms_fusion`` only exposes ``ar_fusion_patterns`` / ``flashinfer_trtllm_fused_allreduce_norm`` when ``flashinfer_comm`` is importable, so the module-level import would raise on CUDA boxes without FlashInfer installed. Wrap the import in ``contextlib.suppress(ImportError)``, default both names to None, and gate every reference behind ``is not None``. The fused-quant impls fall through to their non-flashinfer paths (which use vLLM's own C++ rms-norm kernels), and the no-quant impl falls through to ``tensor_model_parallel_all_reduce`` + ``RMSNorm.forward``. This also fixes the latent issue where the no-quant ``_allreduce_rms_norm`` unconditionally called the flashinfer op on non-CUDA platforms whenever ``needs_ar and x.ndim == 2`` held -- relevant now that every TP model goes through this helper (down_proj / o_proj use reduce_results=False). Signed-off-by: mgoin <mgoin64@gmail.com>
…base
Two changes that travel together:
1. Drop the engine-config lookup for the flashinfer fused-AR-RMS
kernel's ``max_token_num`` arg. The free-function refactor of the
helper pushed ``get_current_vllm_config()`` into the call path, which
runs inside the torch.compile region during ``_dummy_run`` -- where
``set_current_vllm_config`` is not active -- and raised:
AssertionError: Current vLLM config is not set ... model forward
time when config is not set.
Compute the value at the call site instead, using the same math
``FlashInferAllReduce._ensure_workspace`` uses: per-(world_size,
capability) workspace MB cap divided by ``hidden_dim * element_size``.
The MB table is resolved once at import time (its lookup pulls in
lazy imports dynamo can't trace through), and the per-call sizing is
a pure int op. TP=1 didn't trip the original assertion because
``do_allreduce`` was False and the lookup was short-circuited.
2. Pull ``input_quant_key()`` up to ``MMLinearKernel`` (and mirror on
``ScaledMMLinearKernel``) so every linear kernel inherits the
default-None contract. Removes the explicit overrides on
``Fp8BlockScaledMMLinearKernel`` and ``FP8ScaledMMLinearKernel``
(still inherited).
Signed-off-by: mgoin <mgoin64@gmail.com>
The old per-quant-key impls (_allreduce_rms_norm_fp8_static_tensor,
_allreduce_rms_norm_nvfp4, _allreduce_rms_norm) each interleaved a
fused-kernel attempt with an inline naive fallback at the tail. The
decision tree was buried inside each impl; "what path did we actually
take" was hard to read.
Classify the available fusions first, then dispatch. Four explicit,
ordered paths:
A. Single kernel for AR + add + RMSNorm + activation-quant
(flashinfer kARResidualRMSNormFP[48]Quant). Returns QA.
B. Single kernel for AR + (add +) RMSNorm, no quant (flashinfer
kARResidualRMSNorm). Returns plain tensor.
C. Explicit AR (if any) + fused (add +) RMSNorm + activation-quant
(vLLM C++ rms_norm_static_fp8_quant variants). Returns QA.
D. Explicit AR (if any) + plain RMSNorm (vLLM C++ on CUDA,
RMSNorm.forward on other platforms). Returns plain tensor.
Two small dispatch tables -- _FUSED_AR_RMS_QUANT_IMPLS for path A and
_FUSED_RMS_QUANT_IMPLS for path C -- key impls by activation quant key.
QuantizedActivation is returned only when a fused kernel actually
produced the pre-quantised output (A and C); B and D return plain
tensors and let the downstream linear quantise its own input.
Adding a ROCm/AITER backend later is a matter of populating the
dispatch tables -- the dispatcher itself doesn't need to change.
Signed-off-by: mgoin <mgoin64@gmail.com>
The helper's 4th argument was renamed to consumer_linear; update the positional calls in llama / qwen2 / qwen3 to use the keyword form (consumer_linear=...). The final-norm sites passed consumer=None as a keyword and would have raised TypeError under the new signature. Signed-off-by: mgoin <mgoin64@gmail.com>
- Path A (flashinfer AR+RMS+quant) now handles residual=None via the zero-residual + throwaway norm_out trick, matching the compile pass's AllReduceFusedRMSNormStaticQuantFP8Pattern. The AR'd input (left in x) becomes the downstream residual. No current model wiring hits this (do_allreduce always coincides with a real residual today), but it removes the asymmetry with the no-quant path and is forward-looking. - Guard fused_ar_rms_norm_quant with `assert type(norm) is RMSNorm`. The fused kernels apply `gamma * x_normed` directly, so a modified norm (e.g. Gemma's `(1 + gamma)`) would silently produce wrong results. - Assert `x.quant_key == self.input_quant_key()` in the FP8 and NVFP4 apply_weights QA branches so a mismatched QuantizedActivation fails loudly instead of corrupting. - Drop the unquantized op-count test case: the no-quant TP=1 path is a plain RMSNorm passthrough that defers to norm.forward, so it codegens natively under the test's IrOpPriority and no custom op fires to count. Signed-off-by: mgoin <mgoin64@gmail.com>
It was using a two-branch form (full kernel call duplicated per residual case) while the quant impls used inline buffer-selection. Behaviorally identical, just restructured to match: one (kernel_residual, norm_out, out_residual) selection + a single kernel call. The no-quant output is norm_out when allocated (no residual) else x, since this path has no quant_out to return. Signed-off-by: mgoin <mgoin64@gmail.com>
The flashinfer fused-AR workspace is hard-capped in MB (a one-shot FlashInfer limit), so it holds at most MB_cap / (hidden * elem) tokens. The compiler pass guards this per compile-range (is_applicable_for_range: range.end <= max_token_num) and leaves oversize ranges unfused. Manual fusion call sites have no such guard, so the profiling run (max_num_batched_tokens) on a wide-hidden model overran the workspace and tripped the assert (e.g. Llama-4-Scout-FP8: 8192 tokens vs a 6553 token capacity at hidden=5120). Replace the assert with a runtime fallback inside the op: when the batch is wider than the workspace, do a standard all-reduce + the matching C++ rms-norm[+quant] kernels, writing the same output buffers the fused kernel would. The check uses the real num_tokens (the op body is eager / opaque to inductor), so it's per-batch -- fused on small decode steps, unfused only on the oversize prefill/profile -- and needs nothing from VllmConfig. The compiler pass is unaffected (its range guard means it never reaches the fallback). all_reduce is out-of-place, so the fallback runs the norm/quant kernels on the freshly all-reduced tensor and only copies back into allreduce_in when the caller reads it (the no-quant norm result, or the unreached no-residual case). The common FP8/NVFP4 with-residual path is copy-free, and NVFP4 writes its outputs in place via scaled_fp4_quant.out. Handles no-quant, FP8-static, and NVFP4 patterns, honoring the norm_out/residual_out aliasing. Signed-off-by: mgoin <mgoin64@gmail.com>
Baking the manual fusion into the shared Llama/Qwen2 base classes broke tensor parallelism for models that reuse them: - LlamaAttention hardcoded o_proj reduce_results=False, expecting the decoder forward to do the all-reduce. Models reusing LlamaAttention with their own (unfused) decoder forward -- mistral, nemotron_nas (DeciLM), arcee, llama_eagle, llama_eagle3 -- silently dropped the attention all-reduce at TP>1. - Qwen2's final norm called fused_ar_rms_norm_quant(do_allreduce=tp>1) while Qwen2DecoderLayer was never converted (reduce_results=True), so the last layer's output was all-reduced twice (confirmed: garbage output at TP=2 on Qwen2.5-0.5B/1.5B). - aria inherited the converted forward but its MoE mlp has no gate_up_proj -> AttributeError. Fix: make fusion opt-in. - LlamaAttention gains a reduce_results param (default True, safe for reusers); the decoder only sets it False when fusing. - LlamaDecoderLayer gains fuse_allreduce (default True) that gates both reduce_results on o_proj/down_proj and the forward path (fused vs the original plain path). The 4 subclasses that don't fully convert (mistral, llama_eagle, llama_eagle3, aria) pass fuse_allreduce=False, restoring their original behavior (also fixes eagle's Identity-norm crash and aria's gate_up_proj crash). - DeciLM/Arcee build attention directly, so the True default fixes them with no per-model change. - Revert the Qwen2 final-norm call to self.norm(...); drop the now-unused import and tp_size. - Defensive getattr(self.mlp, "gate_up_proj", None) in the fused path. Tested TP=1 vs TP=2 (greedy): Qwen2.5-1.5B now matches (was garbage); Llama-3.2-1B-FP8 fused path unchanged; Mistral-7B coherent (AR restored; its code path is now identical to pre-fusion main). Manual-fusion op-count test still passes. Signed-off-by: mgoin <mgoin64@gmail.com>
…llm-project#43224) Replace the per-call do_allreduce= booleans and the fuse_allreduce decoder flag with a LayerCommunicator that owns all-reduce + residual-add + RMSNorm(+quant) for each decoder layer. The decoder forward becomes: hidden, residual = self.comm.prepare_attn(hidden, residual) hidden = self.self_attn(...) hidden, residual = self.comm.prepare_mlp(hidden, residual) hidden = self.mlp(hidden) (cleaner than the pre-fusion forward: no `residual is None` branch). The communicator reads the layer's norms/linears live and derives whether to reduce from the row-linear's existing reduce_results -- the single source of truth -- so there is no fusion-specific flag. A FULL input never re-reduces, so a double all-reduce is impossible by construction; a non-RMSNorm norm (eagle's Identity layer-0) is tolerated. - Llama / Qwen3 / Mistral fuse by default (no flag in the model def). - Reusers whose linears reduce themselves opt out with reduce_results=True (aria MoE; eagle/eagle3 custom forwards). - Mistral migrates onto the communicator (rare ada variant keeps its explicit path). Custom-forward reusers that consume the final hidden directly (EagleMistral, ErnieMTP) reduce it via finalize_norm / reduce_results=True. Sequence-parallelism is not implemented; Scatter.SHARD + reduce-scatter/ all-gather slot into _reduce_norm without touching any model forward. Verified byte-identical to baseline at TP1 and TP2: Llama-3.1-8B and migrated Mistral-7B (fused == unfused baseline). Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: mgoin <mgoin64@gmail.com>
Tie the manual-fusion QuantizedActivation integration together with two small helpers in quant_activation.py: expose_input_quant_key (the single scheme->layer bridge) and as_quantized_activation (validate + narrow on the consumer side, replacing the duplicated isinstance/assert in the fp8 and nvfp4 kernels). Add GPU-free contract tests pinning which backends consume a pre-quantized activation and that the bridge skips the rest. Signed-off-by: mgoin <mike.goin12@gmail.com> Co-authored-by: Claude Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: mgoin <mgoin64@gmail.com>
# Conflicts: # tests/compile/fusions_e2e/models.py # tests/fusion/test_quant_activation_contract.py # vllm/model_executor/kernels/linear/nvfp4/flashinfer.py # vllm/model_executor/kernels/linear/scaled_mm/ScaledMMLinearKernel.py # vllm/model_executor/layers/fusion/quant_activation.py Signed-off-by: mgoin <mgoin64@gmail.com>
Make ResidualStream (renamed from the prototype LayerCommunicator) the single manual-fusion entry point model code touches, subsuming the standalone fused_ar_rms_norm_quant helper as an internal kernel-dispatch detail. - Rename communicator.py -> residual_stream.py, LayerCommunicator -> ResidualStream, self.comm -> self.residual_stream; reframe docstring around the residual stream rather than SGLang's communicator. - Convert Qwen3 (the last functional-helper holdout) to the prepare_attn / prepare_mlp path, declaring output_scatter=PARTIAL. - Route Qwen2Model.forward's final norm through finalize_norm. This fixes a latent missing-final-all-reduce at TP>1 for Qwen3 (deferred decoders left the last hidden PARTIAL, but the inherited plain self.norm never reduced it). Unmigrated decoders expose no output_scatter -> FULL -> plain norm, identical to the pre-fusion path. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: mgoin <mgoin64@gmail.com>
mgoin
marked this pull request as ready for review
June 26, 2026 16:15
mgoin
requested review from
BoyuanFeng,
ProExpertProg,
patrickvonplaten,
sighingnow,
vadiklyutiy,
youkaichao and
zou3519
as code owners
June 26, 2026 16:15
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>
mgoin
commented
Jun 30, 2026
Comment on lines
+329
to
+332
| # Distribution state this layer leaves its output in (read by the model's | ||
| # final norm). Deferred (fused) decoders skip o_proj/down_proj reduce. | ||
| self.output_scatter = Scatter.FULL if reduce_results else Scatter.PARTIAL | ||
| self.residual_stream = ResidualStream(self, vllm_config=vllm_config) |
Member
Author
There was a problem hiding this comment.
@jeejeelee could we maybe pass in the modules directly to ResidualStream here? I'm a bit uncomfortable passing in self given all the model definitions can use different names. It would be more understandable to identify what is "qkv" or "gate_up" by passing in the modules directly as args
Contributor
|
Hi @mgoin, the pre-commit checks have failed. Please run: uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-filesThen, commit the changes and push to your branch. For future commits, |
Contributor
|
This pull request has merge conflicts that must be resolved before it can be |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Moves all-reduce + residual-add + RMSNorm + activation-quant out of the compiler pattern-matching passes and into an explicit per-layer abstraction that model
forwardcalls directly. Part of RFC #43224; builds on theQuantizedActivationlinear-kernel contract (#44260). Supersedes the prototype in #42597 (this is the ResidualStream-unified rewrite adapted to the landed contract).Approach
ResidualStream(layers/fusion/residual_stream.py) owns the(all-reduce +) residual-add + RMSNorm + optional activation-quantat each sublayer boundary, so a migrated decoder'sforwardis justprepare_attn → attn → prepare_mlp → mlp. Distribution state isScatter.{FULL,PARTIAL}; whether a layer defers its reduce is read from the row-linear's existingreduce_results(no new fusion flag).finalize_normowns the last decoder's deferred all-reduce.ar_rms_quant.pyis the internal kernel-dispatch backend (FlashInfer fused AR+RMSNorm+FP8/NVFP4, with unfused fallbacks for batches past the FlashInfer workspace cap). OnlyResidualStreamcalls it — it's the single surface model code touches.reduce_results=Trueorfinalize_norm). Qwen3-VL's custom forward is fixed to route the final norm throughfinalize_norm.Caveats
Test
pytest tests/fusion/—QuantizedActivationcontract + manual-fusion-fires (TP1).tests/compile/fusions_e2e/{test_tp2_ar_rms,test_tp2_async_tp}.py— TP2 (needs 2 GPUs).Developed with AI assistance (Claude); all changed lines human-reviewed and owned by the submitter.