[Performance][torch.compile]: FlashInfer Attention + quant fusion performance issue with TP=4

[ProExpertProg]

As seen in #27080, attention+quant fusion on 4xB200 with the FlashInfer attn backend performs worse than unfused code. This should be resolved so that we can turn on this fusion by default.

cc @nvpohanh @pavanimajety @zou3519 for visibility

[nvpohanh]

Could you share repro steps? I will ask @elvischenv to investigate. thanks

[ProExpertProg]

#27080 has the flags but I can post the exact commands tomorrow morning.

[nvpohanh]

@ProExpertProg There are so many comments in that issue. Could you point us to the correct comment? Thanks!

[ProExpertProg]

This is the relevant piece:

• common: --kv-cache-dtype=fp8 --no-enable-prefix-caching -O.pass_config.enable_noop=true
• dynamo partition: -O.use_inductor_graph_partition=False -O.cudagraph_mode=FULL_AND_PIECEWISE
• inductor partition: -O.use_inductor_graph_partition=True -O.cudagraph_mode=FULL_AND_PIECEWISE
• no partition: -O.use_inductor_graph_partition=False -O.cudagraph_mode=FULL_DECODE_ONLY -O.splitting_ops=[]
• fusion: -O.pass_config.enable_attn_fusion=true

So, these are the commands:

# inductor partition
vllm serve redhatai/meta-llama-3.1-70B-Instruct-FP8 --tensor-parallel-size 4 --kv-cache-dtype=fp8 --no-enable-prefix-caching -O.pass_config.enable_noop=true -O.use_inductor_graph_partition=True -O.cudagraph_mode=FULL_AND_PIECEWISE

# inductor partition + fusion
vllm serve redhatai/meta-llama-3.1-70B-Instruct-FP8 --tensor-parallel-size 4 --kv-cache-dtype=fp8 --no-enable-prefix-caching -O.pass_config.enable_noop=true -O.use_inductor_graph_partition=True -O.cudagraph_mode=FULL_AND_PIECEWISE -O.pass_config.enable_attn_fusion=true

# no partition
vllm serve redhatai/meta-llama-3.1-70B-Instruct-FP8 --tensor-parallel-size 4 --kv-cache-dtype=fp8 --no-enable-prefix-caching -O.pass_config.enable_noop=true -O.use_inductor_graph_partition=False -O.cudagraph_mode=FULL_DECODE_ONLY -O.splitting_ops=[]

# no partition + fusion
vllm serve redhatai/meta-llama-3.1-70B-Instruct-FP8 --tensor-parallel-size 4 --kv-cache-dtype=fp8 --no-enable-prefix-caching -O.pass_config.enable_noop=true -O.use_inductor_graph_partition=False -O.cudagraph_mode=FULL_DECODE_ONLY -O.splitting_ops=[] -O.pass_config.enable_attn_fusion=true

[nvpohanh]

Assign to @elvischenv

[elvischenv]

I could take a look.