@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / a0e8c74005

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 6, 00:45 UTC

[ROCm]: Update rope+kvcache fusion conditions and disable custom op by default (#36716)

Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>

Rohan Potdar committed Mar 25, 2026 at 15:58 UTC a0e8c74005f5782fad30be912c12cd37fc2813e9

5 files changed +42 -18

docs/design/fusions.md

+1 -1

@@ -22,7 +22,7 @@ or just on the low or high end.

22 | ------------------------------------------------------------------------------ | ---------------------------- | ---------------------------------------------- | ------------------------------ | ------------------ | --------- | ------------ |

23 | [AllReduce + RMSNorm](#allreduce--rmsnorm-fuse_allreduce_rms) | `fuse_allreduce_rms` | All-reduce → RMSNorm (+residual_add) (→ quant) | O2 (Hopper/Blackwell + TP > 1) | 5-20% | No | Low |

25 -| [RoPE + KV-Cache Update](#rope--kv-cache-update-fuse_rope_kvcache) | `fuse_rope_kvcache` | Rotary embedding → KV cache write | O1 (ROCm/AITER only) | TBD | No | Low |

25 +| [RoPE + KV-Cache Update](#rope--kv-cache-update-fuse_rope_kvcache) | `fuse_rope_kvcache` | Rotary embedding → KV cache write | O2 (ROCm/AITER only) | 2-4% | No | Low |

26 | [QK Norm + RoPE](#qk-norm--rope-enable_qk_norm_rope_fusion) | `enable_qk_norm_rope_fusion` | Q/K RMSNorm → rotary embedding | Off by default | 2-3% | No | Low |

docs/design/optimization_levels.md

+3 -1

@@ -56,7 +56,6 @@ Fusions:

56 - `-cc.pass_config.fuse_norm_quant=True`*

57 - `-cc.pass_config.fuse_act_quant=True`*

58 - `-cc.pass_config.fuse_act_padding=True`†

59 -- `-cc.pass_config.fuse_rope_kvcache=True`† (will be moved to O2)

60 \* These fusions are only enabled when either op is using a custom kernel, otherwise Inductor fusion is better.</br>

61 † These fusions are ROCm-only and require AITER.

@@ -71,6 +70,9 @@ Settings (on top of `-O1`):

71 - `-cc.cudagraph_mode=FULL_AND_PIECEWISE`

72 - `-cc.pass_config.fuse_allreduce_rms=True`

73 +- `-cc.pass_config.fuse_rope_kvcache=True`†

74 +

75 +† These fusions are ROCm-only and require AITER.

77 ### `-O3`: Aggressive Optimization

vllm/config/compilation.py

+33 -7

@@ -1042,13 +1042,6 @@ class CompilationConfig:

1042 self.splitting_ops = []

1043 return

1044

1045 - # NOTE: this function needs to be called only when mode is

1046 - # CompilationMode.VLLM_COMPILE

1047 - assert self.mode == CompilationMode.VLLM_COMPILE, (

1048 - "set_splitting_ops_for_v1 should only be called when "

1049 - "mode is CompilationMode.VLLM_COMPILE"

1050 - )

1051 -

1045 if self.pass_config.fuse_attn_quant and not self.use_inductor_graph_partition:

1046 self.set_splitting_ops_for_attn_fusion()

1047 else:

@@ -1069,6 +1062,16 @@ class CompilationConfig:

1062 # that doesn't seem to affect performance.

1063 # https://github.com/vllm-project/vllm/issues/33267

1064 if not self.use_inductor_graph_partition:

1065 + if self.pass_config.fuse_rope_kvcache:

1066 + logger.warning_once(

1067 + "fuse_rope_kvcache is enabled, but splitting_ops is None "

1068 + "and Inductor graph partition is not enabled."

1069 + "Disabling fuse_rope_kvcache."

1070 + "Please either set splitting_ops to an empty list []"

1071 + "or set use_inductor_graph_partition to True "

1072 + "to enable RoPE+KV cache fusion."

1073 + )

1074 + self.pass_config.fuse_rope_kvcache = False

1075 self.splitting_ops.append("vllm::unified_kv_cache_update")

1076 self.splitting_ops.append("vllm::unified_mla_kv_cache_update")

1077

@@ -1142,6 +1145,29 @@ class CompilationConfig:

1145 op in self.splitting_ops for op in self._attention_ops

1146 )

1147

1148 + def splitting_ops_contain_kv_cache_update(self) -> bool:

1149 + # when using Dynamo partition while splitting ops is None

1150 + # and attn+quant fusion disabled, the kv_cache_update_ops are

1151 + # appended to splitting_ops in set_splitting_ops_for_v1 due to

1152 + # https://github.com/vllm-project/vllm/issues/33267

1153 + # In this case, we return True if the kv_cache_update_ops

1154 + # are not in the splitting_ops yet, but will subsequently

1155 + # be added to splitting_ops.

1156 + if (

1157 + not self.use_inductor_graph_partition

1158 + and self.splitting_ops is None

1159 + and not self.pass_config.fuse_attn_quant

1160 + ):

1161 + return True

1162 +

1163 + kv_cache_update_ops = [

1164 + "vllm::unified_kv_cache_update",

1165 + "vllm::unified_mla_kv_cache_update",

1166 + ]

1167 + return self.splitting_ops is not None and all(

1168 + op in self.splitting_ops for op in kv_cache_update_ops

1169 + )

1170 +

1171 def is_attention_compiled_piecewise(self) -> bool:

1172 if not self.splitting_ops_contain_attention():

1173 return False

vllm/config/vllm.py

+5 -2

@@ -144,7 +144,10 @@ def enable_rope_kvcache_fusion(cfg: "VllmConfig") -> bool:

144 return (

145 rocm_aiter_ops.is_enabled()

146 and cfg.compilation_config.is_custom_op_enabled("rotary_embedding")

147 - and cfg.compilation_config.use_inductor_graph_partition

147 + and (

148 + cfg.compilation_config.use_inductor_graph_partition

149 + or not cfg.compilation_config.splitting_ops_contain_kv_cache_update()

150 + )

151 )

152

153

@@ -190,7 +193,7 @@ OPTIMIZATION_LEVEL_01 = {

193 "enable_sp": False,

194 "fuse_gemm_comms": False,

195 "fuse_act_padding": enable_norm_pad_fusion,

193 - "fuse_rope_kvcache": enable_rope_kvcache_fusion,

196 + "fuse_rope_kvcache": False,

197 },

198 "cudagraph_mode": CUDAGraphMode.PIECEWISE,

199 "use_inductor_graph_partition": False,

vllm/platforms/rocm.py

-7

@@ -649,13 +649,6 @@ class RocmPlatform(Platform):

649 and "-grouped_topk" not in compilation_config.custom_ops

650 ):

651 compilation_config.custom_ops.append("+grouped_topk")

652 - # Enable rotary embedding customop when using AITER if not disabled by user

653 - if (

654 - rocm_aiter_ops.is_enabled()

655 - and "+rotary_embedding" not in compilation_config.custom_ops

656 - and "-rotary_embedding" not in compilation_config.custom_ops

657 - ):

658 - compilation_config.custom_ops.append("+rotary_embedding")

652

653 # Default dispatch to rocm's sparse_attn_indexer implementation

654 compilation_config.custom_ops.append("+sparse_attn_indexer")