[RFC]: Changes to support attention + quant fusion

[ProExpertProg]

Motivation.

I am working on fusing FP8 quantization onto (V0) attention with a new torch.compile pass. However, this pass will require some other changes to core vLLM code so I just wanted to collect people's opinions before proceeding.

AMD saw manually fusing quantization into the attention op lead to a major improvement in performance in the ROCm fork. Basically, when fused, the attention output is quantized inside the attention Triton kernel before being written out, saving data movement by avoiding a full roundtrip to global memory at 16-bit precision.

To avoid modifying the model definitions (approach used in ROCm fork), we are planning to do the fusion using a torch.compile pass. Apart from the pass implementation, there are two obstacles in the Python integration code:

1. Attention is wrapped into a unified_attention custom op
2. In V1, piecewise compilation means that the ops we want to fuse appear in different graphs

Proposed Change.

1. Extend unified_attention

EDIT: @youkaichao has suggested to directly store output_scale on the attention object, which will also help more easily figure out if the attention backend supports fusing the op. So the proposal below is not necessary unless the approach does not work.

OLD:
For the first issue, I propose adding two optional parameters to the unified_attention op:

• output_scale: Optional[torch.Tensor] = None
• quant_config: Optional[QuantConfig] = None, which contains:
  • dynamic: bool (is the quantization dynamic or static)
  • per_token: bool (is quant per-token or per-tensor - this could later be extended to other quantization schemes (group, etc.) if necessary)
  • dtype: torch.dtype
  • This could also be named o_quant_config or out_quant_config not to be confused with quantization config for QKV and P.

This way, we can make the following fx.Graph replacement:

# Original
output = torch.ops.vllm.unified_attention(q, k, v, 'model.layers.0.self_attn.attn')
output_quant = torch.empty(...)
torch.ops._C.static_scaled_fp8_quant(output_quant, output, scale) # in-place op

# Fused 
output_quant = torch.ops.vllm.unified_attention(q, k, v, 'model.layers.0.self_attn.attn', output_scale=scale, quant_config=(True, True, torch.fp8_e4m3fn))

unified_attention_with_output (often used in V1) can use the same approach with slightly different handling for output (and a separate pattern matcher pattern). If we only used unified_attention_with_output, quant_config.dtype is redundant with the passed output dtype but we need to support both and redundancy shouldn't hurt us here.

2. Add "sticky ops" for graph splitting

Once we want to add support for attention+quant fusion to V1, we'd have to revisit our approach to splitting graphs into piecewise: basically, we need to make sure that scaled_fp8_quant ends in the same graph as attention. For that, I suggest introducing the concept of "sticky ops", whether that's by name or via callback, which are ops that use the output of attention and end up in the attention sub-graph. We can only enable this when fusion is enabled & possible to make sure we still benefit from CUDA graphs for these ops when not fusing.

Finally, this will further cement our reliance on pattern matching in-place ops. So we should figure out #14703.

Feedback Period.

This week (April 7-11)

CC List.

@youkaichao @WoosukKwon @tlrmchlsmth @SageMoore @LucasWilkinson @rasmith @gshtras

Any Other Things.

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[rasmith]

Maybe call it input_scale since that coincides with the actual parameter name the values are coming from.

[ProExpertProg]

Maybe call it input_scale since that coincides with the actual parameter name the values are coming from.

I think that might be confusing just looking at a call to unified_attention as it would imply the scale is for the inputs to the op.

[youkaichao]

output_scale=scale, quant_config=(True, True, torch.fp8_e4m3fn)

if these are static (not changing for every batch of input), you can store it in the attention class.

[ProExpertProg]

output_scale=scale, quant_config=(True, True, torch.fp8_e4m3fn)

if these are static (not changing for every batch of input), you can store it in the attention class.

I guess that would work! So the torch.compile pass modifies the backend/impl instance directly by using the layer key (same lookup that unified_attention does in its implementation?

[ProExpertProg]

I guess how would we access the forward_context during torch.compile, is it available?

[youkaichao]

attention layers are stored in vllm_config.compilation_config.static_forward_context

[rasmith]

Maybe call it input_scale since that coincides with the actual parameter name the values are coming from.

I think that might be confusing just looking at a call to unified_attention as it would imply the scale is for the inputs to the op.

Well, I say this because there is already an output_scale parameter which could also be confusing.

[rasmith]

Why store in the attention object when the value is already loaded into o_proj.input_scale which is RowParallelLinear?

[ProExpertProg]

Why store in the attention object when the value is already loaded into o_proj.input_scale which is RowParallelLinear?

The attention object doesn't have access to o_proj

[zou3519]

cc @drisspg @BoyuanFeng @eellison (the CUDAGraphs Piecewise refactor might become a moving target, though we don't need to discuss that in this thread).

[ProExpertProg]

@youkaichao I'm realizing that returning both 16-bit (before fusion) and 8-bit values (after quant fusion) from unified_attention might pose trouble as the fake op (used in tracing) will be wrong some of the time (it currently always returns 16-bit). This will mess with the pattern matcher as it traces the pattern and replacement to figure out data types.

For now, I can just proceed with a manual replacement but this feels like an inconsistency we should address. Additionally, for dynamic quantization, the scale will have to be either passed in or returned so we'll likely have to do something there (we could add a lookup op that grabs the scale from the attention object). But at least dynamic quant is not a priority right now.

In the future, we could either add the dtype as a param (and maybe other quant info), or we could always use unified_attention_with_output - is there a reason we don't always use that interface?

[youkaichao]

In the future, we could either add the dtype as a param (and maybe other quant info), or we could always use unified_attention_with_output - is there a reason we don't always use that interface?

either looks good to me. unified_attention is for some attention kernels that allocate the output internally and do not accept an out buffer. that has some implication on cudagraph buffer usage. But in general I assume unified_attention_with_output would become dominant in the future. For attention kernels that do not take an out buffer, we can use one additional copy to make it compatible with unified_attention_with_output

[ProExpertProg]

@youkaichao I think we don't have access to the scale object in the fx.Graph (it's a parameter to the graph) so I think we should pass it to unified_attention_with_output. The quant metadata can be set on the attention object