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

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 8, 15:15 UTC

[torch.compile][Fusion] Fix attention fusion pass removing kv_udpate op. (#33945)

Signed-off-by: charlifu <charlifu@amd.com>

Charlie Fu committed Feb 9, 2026 at 15:15 UTC bb9f97308d0b88c5ad2d64c217b22866e40c79df

2 files changed +22 -1

tests/compile/passes/test_fusion_attn.py

+12 -1

@@ -267,7 +267,7 @@ elif current_platform.is_rocm():

267 PATTERN_TEST_MODELS_FP8 = [

268 ("amd/Llama-3.1-8B-Instruct-FP8-KV", TestAttentionFp8StaticQuantPatternModel)

269 ]

270 - BACKENDS = [

270 + BACKENDS_FP8 = [

271 AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN,

272 AttentionBackendEnum.ROCM_ATTN,

273 AttentionBackendEnum.TRITON_ATTN,

@@ -474,6 +474,17 @@ def test_attention_quant_pattern(

474 assert attn_nodes_pre[0].kwargs.get("output_block_scale") is None, (

475 "Attention should not have output_block_scale before fusion"

476 )

477 +

478 + kv_cache_dummy_dep_pre_is_none = (

479 + attn_nodes_pre[0].kwargs.get("kv_cache_dummy_dep") is None

480 + )

481 + kv_cache_dummy_dep_post_is_none = (

482 + attn_nodes_post[0].kwargs.get("kv_cache_dummy_dep") is None

483 + )

484 + assert not (kv_cache_dummy_dep_pre_is_none ^ kv_cache_dummy_dep_post_is_none), (

485 + "The kv_cache_dummy_dep should be consistent before and after fusion"

486 + )

487 +

488 if quant_key.dtype == FP8_DTYPE:

489 assert attn_nodes_post[0].kwargs.get("output_block_scale") is None, (

490 "Attention should not have output_block_scale after FP8 fusion"

vllm/compilation/passes/fusion/attn_quant_fusion.py

+10

@@ -142,6 +142,7 @@ class AttentionFp8StaticQuantPattern(AttentionQuantPattern):

142 v: torch.Tensor,

143 output_attn: torch.Tensor,

144 scale: torch.Tensor,

145 + kv_cache_dummy_dep: torch.Tensor,

146 ) -> torch.Tensor:

147 at1 = auto_functionalized(

148 ATTN_OP,

@@ -152,6 +153,7 @@ class AttentionFp8StaticQuantPattern(AttentionQuantPattern):

153 layer_name=self.layer_name,

154 output_scale=None,

155 output_block_scale=None,

156 + kv_cache_dummy_dep=kv_cache_dummy_dep,

157 )

158 attn_out_view = RESHAPE_OP(

159 at1[1], [q.shape[0], self.num_heads * self.head_size]

@@ -165,6 +167,7 @@ class AttentionFp8StaticQuantPattern(AttentionQuantPattern):

167 v: torch.Tensor,

168 output_attn: torch.Tensor,

169 scale: torch.Tensor,

170 + kv_cache_dummy_dep: torch.Tensor,

171 ) -> torch.Tensor:

172 # attn output in quant_dtype

173 output_attn = torch.ops.aten.full.default(

@@ -182,6 +185,7 @@ class AttentionFp8StaticQuantPattern(AttentionQuantPattern):

185 layer_name=self.layer_name,

186 output_scale=scale,

187 output_block_scale=None,

188 + kv_cache_dummy_dep=kv_cache_dummy_dep,

189 )

190 return RESHAPE_OP(at1[1], [-1, self.num_heads * self.head_size])

191

@@ -191,6 +195,7 @@ class AttentionFp8StaticQuantPattern(AttentionQuantPattern):

195 self.empty(5, self.num_heads, self.head_size), # v

196 self.empty(5, self.num_heads, self.head_size), # attn_output

197 empty_fp32(1, 1), # scale

198 + self.empty(0), # kv_cache_dummy_dep

199 ]

200

201 pm.register_replacement(

@@ -228,6 +233,7 @@ class AttentionNvfp4QuantPattern(AttentionQuantPattern):

233 output_quant: torch.Tensor,

234 output_scale: torch.Tensor,

235 input_scale: torch.Tensor,

236 + kv_cache_dummy_dep: torch.Tensor,

237 ) -> tuple[torch.Tensor, torch.Tensor]:

238 at1 = auto_functionalized(

239 ATTN_OP,

@@ -238,6 +244,7 @@ class AttentionNvfp4QuantPattern(AttentionQuantPattern):

244 layer_name=self.layer_name,

245 output_scale=None,

246 output_block_scale=None,

247 + kv_cache_dummy_dep=kv_cache_dummy_dep,

248 )

249 attn_out_view = RESHAPE_OP(

250 at1[1], [q.shape[0], self.num_heads * self.head_size]

@@ -261,6 +268,7 @@ class AttentionNvfp4QuantPattern(AttentionQuantPattern):

268 output_quant: torch.Tensor,

269 output_scale: torch.Tensor,

270 input_scale: torch.Tensor,

271 + kv_cache_dummy_dep: torch.Tensor,

272 ) -> tuple[torch.Tensor, torch.Tensor]:

273 # attention output in quant_dtype

274 output_attn = torch.ops.aten.full.default(

@@ -280,6 +288,7 @@ class AttentionNvfp4QuantPattern(AttentionQuantPattern):

288 layer_name=self.layer_name,

289 output_scale=input_scale,

290 output_block_scale=output_scale_view,

291 + kv_cache_dummy_dep=kv_cache_dummy_dep,

292 )

293 output = RESHAPE_OP(at2[1], [-1, self.num_heads * self.head_size // 2])

294 return output, at2[2]

@@ -294,6 +303,7 @@ class AttentionNvfp4QuantPattern(AttentionQuantPattern):

303 128, round_up(self.num_heads * self.head_size // 16, 4)

304 ), # output_scale

305 empty_fp32(1, 1), # input_scale

306 + self.empty(0), # kv_cache_dummy_dep

307 ]

308

309 pm.register_replacement(