@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / f44afef6d6
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 5, 07:15 UTC

[compile] Allow strings in custom ops without regressing compilation times (#38123)

Signed-off-by: Richard Zou <zou3519@gmail.com>

Richard Zou committed Apr 10, 2026 at 09:26 UTC f44afef6d61faa40be65c14800307f42fa64ca55
16 files changed +565 -168
tests/compile/passes/test_rope_kvcache_fusion.py
+2 -1
@@ -28,6 +28,7 @@ from vllm.forward_context import get_forward_context, set_forward_context
28 from vllm.model_executor.layers.attention import Attention
29 from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
30 from vllm.platforms import current_platform
31 +from vllm.utils.torch_utils import _encode_layer_name
32 from vllm.v1.attention.backend import (
33 AttentionBackend,
34 CommonAttentionMetadata,
@@ -170,7 +171,7 @@ class QKRoPEKVCacheTestModel(torch.nn.Module):
171 k = k.view(-1, self.num_kv_heads, self.head_size)
172 v = v.view(-1, self.num_kv_heads, self.head_size)
173 kv_cache_dummy_dep = torch.ops.vllm.unified_kv_cache_update(
173 - k, v, self.layer_name
174 + k, v, _encode_layer_name(self.layer_name)
175 )
176 return q, k, v, kv_cache_dummy_dep
177
vllm/compilation/passes/fusion/attn_quant_fusion.py
+182 -50
@@ -17,6 +17,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
17 )
18 from vllm.platforms import current_platform
19 from vllm.utils.math_utils import round_up
20 +from vllm.utils.torch_utils import _USE_LAYERNAME, _encode_layer_name
21
22 from ..vllm_inductor_pass import VllmFusionPatternMatcherPass, VllmPatternReplacement
23 from .matcher_utils import MatcherQuantFP8
@@ -53,21 +54,43 @@ class AttnFp8StaticQuantPattern(VllmPatternReplacement[..., torch.Tensor]):
54
55 @property
56 def pattern(self) -> Callable[..., torch.Tensor]:
56 - def _pattern(
57 - q: torch.Tensor,
58 - k: torch.Tensor,
59 - v: torch.Tensor,
60 - output_attn: torch.Tensor,
61 - scale: torch.Tensor,
62 - kv_cache_dummy_dep: torch.Tensor,
63 - ) -> torch.Tensor:
57 + # When _USE_LAYERNAME is enabled (torch >= 2.11), layer_name is
58 + # passed as an explicit pattern input so the pattern matcher
59 + # treats it as a wildcard matching hoisted LayerName placeholders.
60 + # Otherwise it stays as a closure constant (original behavior).
61 + _ln = _encode_layer_name(self._layer_name)
62 +
63 + if _USE_LAYERNAME:
64 +
65 + def _pattern_with_ln( # type: ignore[misc]
66 + q, k, v, output_attn, scale, kv_cache_dummy_dep, layer_name
67 + ):
68 + at1 = auto_functionalized(
69 + ATTN_OP,
70 + query=q,
71 + key=k,
72 + value=v,
73 + output=output_attn,
74 + layer_name=layer_name,
75 + output_scale=None,
76 + output_block_scale=None,
77 + kv_cache_dummy_dep=kv_cache_dummy_dep,
78 + )
79 + attn_out_view = RESHAPE_OP(
80 + at1[1], [q.shape[0], self._num_heads * self._head_size]
81 + )
82 + return self._quant_matcher(attn_out_view, scale)[0]
83 +
84 + return _pattern_with_ln
85 +
86 + def _pattern(q, k, v, output_attn, scale, kv_cache_dummy_dep):
87 at1 = auto_functionalized(
88 ATTN_OP,
89 query=q,
90 key=k,
91 value=v,
92 output=output_attn,
70 - layer_name=self._layer_name,
93 + layer_name=_ln,
94 output_scale=None,
95 output_block_scale=None,
96 kv_cache_dummy_dep=kv_cache_dummy_dep,
@@ -81,14 +104,34 @@ class AttnFp8StaticQuantPattern(VllmPatternReplacement[..., torch.Tensor]):
104
105 @property
106 def replacement(self) -> Callable[..., torch.Tensor]:
84 - def _replacement(
85 - q: torch.Tensor,
86 - k: torch.Tensor,
87 - v: torch.Tensor,
88 - output_attn: torch.Tensor,
89 - scale: torch.Tensor,
90 - kv_cache_dummy_dep: torch.Tensor,
91 - ) -> torch.Tensor:
107 + _ln = _encode_layer_name(self._layer_name)
108 +
109 + if _USE_LAYERNAME:
110 +
111 + def _replacement_with_ln( # type: ignore[misc]
112 + q, k, v, output_attn, scale, kv_cache_dummy_dep, layer_name
113 + ):
114 + output_attn = torch.empty(
115 + [q.shape[0], self._num_heads, self._head_size],
116 + dtype=FP8_DTYPE,
117 + device=q.device,
118 + )
119 + at1 = auto_functionalized(
120 + ATTN_OP,
121 + query=q,
122 + key=k,
123 + value=v,
124 + output=output_attn,
125 + layer_name=layer_name,
126 + output_scale=scale,
127 + output_block_scale=None,
128 + kv_cache_dummy_dep=kv_cache_dummy_dep,
129 + )
130 + return RESHAPE_OP(at1[1], [-1, self._num_heads * self._head_size])
131 +
132 + return _replacement_with_ln
133 +
134 + def _replacement(q, k, v, output_attn, scale, kv_cache_dummy_dep):
135 output_attn = torch.empty(
136 [q.shape[0], self._num_heads, self._head_size],
137 dtype=FP8_DTYPE,
@@ -100,7 +143,7 @@ class AttnFp8StaticQuantPattern(VllmPatternReplacement[..., torch.Tensor]):
143 key=k,
144 value=v,
145 output=output_attn,
103 - layer_name=self._layer_name,
146 + layer_name=_ln,
147 output_scale=scale,
148 output_block_scale=None,
149 kv_cache_dummy_dep=kv_cache_dummy_dep,
@@ -113,7 +156,7 @@ class AttnFp8StaticQuantPattern(VllmPatternReplacement[..., torch.Tensor]):
156 dtype = self._dtype
157 num_heads = self._num_heads
158 head_size = self._head_size
116 - return [
159 + inputs: list = [
160 self.empty(5, num_heads, head_size, dtype=dtype), # q
161 self.empty(5, num_heads, head_size, dtype=dtype), # k
162 self.empty(5, num_heads, head_size, dtype=dtype), # v
@@ -121,6 +164,9 @@ class AttnFp8StaticQuantPattern(VllmPatternReplacement[..., torch.Tensor]):
164 self.empty_fp32(1, 1), # scale
165 self.empty(0, dtype=dtype), # kv_cache_dummy_dep
166 ]
167 + if _USE_LAYERNAME:
168 + inputs.append(_encode_layer_name(self._layer_name))
169 + return inputs
170
171
172 class AttnNvfp4QuantPattern(
@@ -144,23 +190,64 @@ class AttnNvfp4QuantPattern(
190
191 @property
192 def pattern(self) -> Callable[..., tuple[torch.Tensor, torch.Tensor]]:
193 + _ln = _encode_layer_name(self._layer_name)
194 +
195 + if _USE_LAYERNAME:
196 +
197 + def _pattern_with_ln( # type: ignore[misc]
198 + q,
199 + k,
200 + v,
201 + output_attn,
202 + output_quant,
203 + output_scale,
204 + input_scale,
205 + kv_cache_dummy_dep,
206 + layer_name,
207 + ):
208 + at1 = auto_functionalized(
209 + ATTN_OP,
210 + query=q,
211 + key=k,
212 + value=v,
213 + output=output_attn,
214 + layer_name=layer_name,
215 + output_scale=None,
216 + output_block_scale=None,
217 + kv_cache_dummy_dep=kv_cache_dummy_dep,
218 + )
219 + attn_out_view = RESHAPE_OP(
220 + at1[1], [q.shape[0], self._num_heads * self._head_size]
221 + )
222 + at2 = auto_functionalized(
223 + self._QUANT_OP,
224 + input=attn_out_view,
225 + input_scale=input_scale,
226 + is_sf_swizzled_layout=True,
227 + output=output_quant,
228 + output_scale=output_scale,
229 + )
230 + return at2[1], torch.ops.aten.view.dtype(at2[2], FP8_DTYPE)
231 +
232 + return _pattern_with_ln
233 +
234 def _pattern(
148 - q: torch.Tensor,
149 - k: torch.Tensor,
150 - v: torch.Tensor,
151 - output_attn: torch.Tensor,
152 - output_quant: torch.Tensor,
153 - output_scale: torch.Tensor,
154 - input_scale: torch.Tensor,
155 - kv_cache_dummy_dep: torch.Tensor,
156 - ) -> tuple[torch.Tensor, torch.Tensor]:
235 + q,
236 + k,
237 + v,
238 + output_attn,
239 + output_quant,
240 + output_scale,
241 + input_scale,
242 + kv_cache_dummy_dep,
243 + ):
244 at1 = auto_functionalized(
245 ATTN_OP,
246 query=q,
247 key=k,
248 value=v,
249 output=output_attn,
163 - layer_name=self._layer_name,
250 + layer_name=_ln,
251 output_scale=None,
252 output_block_scale=None,
253 kv_cache_dummy_dep=kv_cache_dummy_dep,
@@ -176,42 +263,80 @@ class AttnNvfp4QuantPattern(
263 output=output_quant,
264 output_scale=output_scale,
265 )
179 - output_scale_view = torch.ops.aten.view.dtype(at2[2], FP8_DTYPE)
180 - return at2[1], output_scale_view
266 + return at2[1], torch.ops.aten.view.dtype(at2[2], FP8_DTYPE)
267
268 return _pattern
269
270 @property
271 def replacement(self) -> Callable[..., tuple[torch.Tensor, torch.Tensor]]:
272 + _ln = _encode_layer_name(self._layer_name)
273 +
274 + if _USE_LAYERNAME:
275 +
276 + def _replacement_with_ln( # type: ignore[misc]
277 + q,
278 + k,
279 + v,
280 + output_attn,
281 + _output_quant,
282 + output_scale,
283 + input_scale,
284 + kv_cache_dummy_dep,
285 + layer_name,
286 + ):
287 + output_attn = torch.empty(
288 + [q.shape[0], self._num_heads, self._head_size // 2],
289 + dtype=FP4_DTYPE,
290 + device=q.device,
291 + )
292 + osv = torch.ops.aten.view.dtype(output_scale, FP8_DTYPE)
293 + at2 = auto_functionalized(
294 + ATTN_OP,
295 + query=q,
296 + key=k,
297 + value=v,
298 + output=output_attn,
299 + layer_name=layer_name,
300 + output_scale=input_scale,
301 + output_block_scale=osv,
302 + kv_cache_dummy_dep=kv_cache_dummy_dep,
303 + )
304 + return RESHAPE_OP(
305 + at2[1], [-1, self._num_heads * self._head_size // 2]
306 + ), at2[2]
307 +
308 + return _replacement_with_ln
309 +
310 def _replacement(
187 - q: torch.Tensor,
188 - k: torch.Tensor,
189 - v: torch.Tensor,
190 - output_attn: torch.Tensor,
191 - _output_quant: torch.Tensor,
192 - output_scale: torch.Tensor,
193 - input_scale: torch.Tensor,
194 - kv_cache_dummy_dep: torch.Tensor,
195 - ) -> tuple[torch.Tensor, torch.Tensor]:
311 + q,
312 + k,
313 + v,
314 + output_attn,
315 + _output_quant,
316 + output_scale,
317 + input_scale,
318 + kv_cache_dummy_dep,
319 + ):
320 output_attn = torch.empty(
321 [q.shape[0], self._num_heads, self._head_size // 2],
322 dtype=FP4_DTYPE,
323 device=q.device,
324 )
201 - output_scale_view = torch.ops.aten.view.dtype(output_scale, FP8_DTYPE)
325 + osv = torch.ops.aten.view.dtype(output_scale, FP8_DTYPE)
326 at2 = auto_functionalized(
327 ATTN_OP,
328 query=q,
329 key=k,
330 value=v,
331 output=output_attn,
208 - layer_name=self._layer_name,
332 + layer_name=_ln,
333 output_scale=input_scale,
210 - output_block_scale=output_scale_view,
334 + output_block_scale=osv,
335 kv_cache_dummy_dep=kv_cache_dummy_dep,
336 )
213 - output = RESHAPE_OP(at2[1], [-1, self._num_heads * self._head_size // 2])
214 - return output, at2[2]
337 + return RESHAPE_OP(
338 + at2[1], [-1, self._num_heads * self._head_size // 2]
339 + ), at2[2]
340
341 return _replacement
342
@@ -219,18 +344,19 @@ class AttnNvfp4QuantPattern(
344 dtype = self._dtype
345 num_heads = self._num_heads
346 head_size = self._head_size
222 - return [
347 + inputs: list = [
348 self.empty_bf16(5, num_heads, head_size), # q
349 self.empty_bf16(5, num_heads, head_size), # k
350 self.empty_bf16(5, num_heads, head_size), # v
351 self.empty_bf16(5, num_heads, head_size), # output_attn
227 - self.empty(5, num_heads * head_size // 2, dtype=FP4_DTYPE), # output_quant
228 - self.empty_i32(
229 - 128, round_up(num_heads * head_size // 16, 4)
230 - ), # output_scale
352 + self.empty(5, num_heads * head_size // 2, dtype=FP4_DTYPE),
353 + self.empty_i32(128, round_up(num_heads * head_size // 16, 4)),
354 self.empty_fp32(1, 1), # input_scale
355 self.empty(0, dtype=dtype), # kv_cache_dummy_dep
356 ]
357 + if _USE_LAYERNAME:
358 + inputs.append(_encode_layer_name(self._layer_name))
359 + return inputs
360
361
362 class AttnQuantFusionPass(VllmFusionPatternMatcherPass):
@@ -259,13 +385,19 @@ class AttnQuantFusionPass(VllmFusionPatternMatcherPass):
385 "so no fusion patterns were registered."
386 )
387
388 + # When _USE_LAYERNAME is enabled, layer_name is a wildcard so all
389 + # layers produce the same pattern — register once then break.
390 for layer in layers:
391 if layer.impl.fused_output_quant_supported(_FP8_QUANT_KEY):
392 self.register(AttnFp8StaticQuantPattern(layer, dtype))
393 + if _USE_LAYERNAME:
394 + break
395
396 if current_platform.is_cuda() and hasattr(torch.ops._C, "scaled_fp4_quant"):
397 for layer in layers:
398 if layer.impl.fused_output_quant_supported(kNvfp4Dynamic):
399 self.register(AttnNvfp4QuantPattern(layer, dtype))
400 + if _USE_LAYERNAME:
401 + break
402
403 self.dump_patterns(config, self.pm_pass)
vllm/compilation/passes/fusion/mla_attn_quant_fusion.py
+166 -36
@@ -15,6 +15,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
15 kNvfp4Dynamic,
16 )
17 from vllm.platforms import current_platform
18 +from vllm.utils.torch_utils import _USE_LAYERNAME, _encode_layer_name
19
20 from ..vllm_inductor_pass import VllmFusionPatternMatcherPass, VllmPatternReplacement
21 from .matcher_utils import MatcherQuantFP8
@@ -49,21 +50,43 @@ class MLAAttnFp8StaticQuantPattern(VllmPatternReplacement[..., torch.Tensor]):
50
51 @property
52 def pattern(self) -> Callable[..., torch.Tensor]:
52 - def _pattern(
53 - q: torch.Tensor,
54 - kv_c_normed: torch.Tensor,
55 - k_pe: torch.Tensor,
56 - output_attn: torch.Tensor,
57 - scale: torch.Tensor,
58 - kv_cache_dummy_dep: torch.Tensor,
59 - ) -> torch.Tensor:
53 + _ln = _encode_layer_name(self._layer_name)
54 +
55 + if _USE_LAYERNAME:
56 +
57 + def _pattern_with_ln( # type: ignore[misc]
58 + q,
59 + kv_c_normed,
60 + k_pe,
61 + output_attn,
62 + scale,
63 + kv_cache_dummy_dep,
64 + layer_name,
65 + ):
66 + at1 = auto_functionalized(
67 + MLA_ATTN_OP,
68 + q=q,
69 + kv_c_normed=kv_c_normed,
70 + k_pe=k_pe,
71 + output=output_attn,
72 + layer_name=layer_name,
73 + output_scale=None,
74 + output_block_scale=None,
75 + kv_cache_dummy_dep=kv_cache_dummy_dep,
76 + )
77 + # MLA output is already 2D (T, N*V), no reshape needed
78 + return self._quant_matcher(at1[1], scale)[0]
79 +
80 + return _pattern_with_ln
81 +
82 + def _pattern(q, kv_c_normed, k_pe, output_attn, scale, kv_cache_dummy_dep):
83 at1 = auto_functionalized(
84 MLA_ATTN_OP,
85 q=q,
86 kv_c_normed=kv_c_normed,
87 k_pe=k_pe,
88 output=output_attn,
66 - layer_name=self._layer_name,
89 + layer_name=_ln,
90 output_scale=None,
91 output_block_scale=None,
92 kv_cache_dummy_dep=kv_cache_dummy_dep,
@@ -75,14 +98,41 @@ class MLAAttnFp8StaticQuantPattern(VllmPatternReplacement[..., torch.Tensor]):
98
99 @property
100 def replacement(self) -> Callable[..., torch.Tensor]:
78 - def _replacement(
79 - q: torch.Tensor,
80 - kv_c_normed: torch.Tensor,
81 - k_pe: torch.Tensor,
82 - output_attn: torch.Tensor,
83 - scale: torch.Tensor,
84 - kv_cache_dummy_dep: torch.Tensor,
85 - ) -> torch.Tensor:
101 + _ln = _encode_layer_name(self._layer_name)
102 +
103 + if _USE_LAYERNAME:
104 +
105 + def _replacement_with_ln( # type: ignore[misc]
106 + q,
107 + kv_c_normed,
108 + k_pe,
109 + output_attn,
110 + scale,
111 + kv_cache_dummy_dep,
112 + layer_name,
113 + ):
114 + # MLA output in quant_dtype
115 + output_attn = torch.empty(
116 + [q.shape[0], self._output_dim],
117 + dtype=FP8_DTYPE,
118 + device=q.device,
119 + )
120 + at1 = auto_functionalized(
121 + MLA_ATTN_OP,
122 + q=q,
123 + kv_c_normed=kv_c_normed,
124 + k_pe=k_pe,
125 + output=output_attn,
126 + layer_name=layer_name,
127 + output_scale=scale,
128 + output_block_scale=None,
129 + kv_cache_dummy_dep=kv_cache_dummy_dep,
130 + )
131 + return at1[1]
132 +
133 + return _replacement_with_ln
134 +
135 + def _replacement(q, kv_c_normed, k_pe, output_attn, scale, kv_cache_dummy_dep):
136 # MLA output in quant_dtype
137 output_attn = torch.empty(
138 [q.shape[0], self._output_dim],
@@ -95,7 +145,7 @@ class MLAAttnFp8StaticQuantPattern(VllmPatternReplacement[..., torch.Tensor]):
145 kv_c_normed=kv_c_normed,
146 k_pe=k_pe,
147 output=output_attn,
98 - layer_name=self._layer_name,
148 + layer_name=_ln,
149 output_scale=scale,
150 output_block_scale=None,
151 kv_cache_dummy_dep=kv_cache_dummy_dep,
@@ -105,7 +155,7 @@ class MLAAttnFp8StaticQuantPattern(VllmPatternReplacement[..., torch.Tensor]):
155 return _replacement
156
157 def get_inputs(self) -> list[torch.Tensor]:
108 - return [
158 + inputs: list = [
159 self.empty(5, self._num_heads, self._qk_head_dim, dtype=self._dtype),
160 self.empty(5, self._kv_lora_rank, dtype=self._dtype),
161 self.empty(5, 1, self._qk_rope_head_dim, dtype=self._dtype),
@@ -113,6 +163,9 @@ class MLAAttnFp8StaticQuantPattern(VllmPatternReplacement[..., torch.Tensor]):
163 self.empty_fp32(1, 1),
164 self.empty(0, dtype=self._dtype),
165 ]
166 + if _USE_LAYERNAME:
167 + inputs.append(_encode_layer_name(self._layer_name))
168 + return inputs
169
170
171 class MLAAttnNvfp4QuantPattern(
@@ -141,21 +194,56 @@ class MLAAttnNvfp4QuantPattern(
194 def pattern(
195 self,
196 ) -> Callable[..., tuple[torch.Tensor, torch.Tensor]]:
197 + _ln = _encode_layer_name(self._layer_name)
198 +
199 + if _USE_LAYERNAME:
200 +
201 + def _pattern_with_ln( # type: ignore[misc]
202 + q,
203 + kv_c_normed,
204 + k_pe,
205 + output_attn,
206 + input_scale,
207 + kv_cache_dummy_dep,
208 + layer_name,
209 + ):
210 + at1 = auto_functionalized(
211 + MLA_ATTN_OP,
212 + q=q,
213 + kv_c_normed=kv_c_normed,
214 + k_pe=k_pe,
215 + output=output_attn,
216 + layer_name=layer_name,
217 + output_scale=None,
218 + output_block_scale=None,
219 + kv_cache_dummy_dep=kv_cache_dummy_dep,
220 + )
221 + output_quant, output_scale = create_fp4_output_tensors(
222 + at1[1].shape[0], at1[1].shape[1], at1[1].device, True
223 + )
224 + at2 = auto_functionalized(
225 + self._QUANT_OP,
226 + input=at1[1],
227 + input_scale=input_scale,
228 + is_sf_swizzled_layout=True,
229 + output=output_quant,
230 + output_scale=output_scale,
231 + )
232 + output_scale_view = torch.ops.aten.view.dtype(at2[2], FP8_DTYPE)
233 + return at2[1], output_scale_view
234 +
235 + return _pattern_with_ln
236 +
237 def _pattern(
145 - q: torch.Tensor,
146 - kv_c_normed: torch.Tensor,
147 - k_pe: torch.Tensor,
148 - output_attn: torch.Tensor,
149 - input_scale: torch.Tensor,
150 - kv_cache_dummy_dep: torch.Tensor,
151 - ) -> tuple[torch.Tensor, torch.Tensor]:
238 + q, kv_c_normed, k_pe, output_attn, input_scale, kv_cache_dummy_dep
239 + ):
240 at1 = auto_functionalized(
241 MLA_ATTN_OP,
242 q=q,
243 kv_c_normed=kv_c_normed,
244 k_pe=k_pe,
245 output=output_attn,
158 - layer_name=self._layer_name,
246 + layer_name=_ln,
247 output_scale=None,
248 output_block_scale=None,
249 kv_cache_dummy_dep=kv_cache_dummy_dep,
@@ -182,14 +270,47 @@ class MLAAttnNvfp4QuantPattern(
270 def replacement(
271 self,
272 ) -> Callable[..., tuple[torch.Tensor, torch.Tensor]]:
273 + _ln = _encode_layer_name(self._layer_name)
274 +
275 + if _USE_LAYERNAME:
276 +
277 + def _replacement_with_ln( # type: ignore[misc]
278 + q,
279 + kv_c_normed,
280 + k_pe,
281 + output_attn,
282 + input_scale,
283 + kv_cache_dummy_dep,
284 + layer_name,
285 + ):
286 + # MLA output in quant_dtype (FP4 packed as uint8)
287 + output_attn = torch.empty(
288 + [q.shape[0], self._output_dim // 2],
289 + dtype=FP4_DTYPE,
290 + device=q.device,
291 + )
292 + output_scale = create_fp4_output_tensors(
293 + q.shape[0], self._output_dim, q.device, True
294 + )[1]
295 + output_scale_view = torch.ops.aten.view.dtype(output_scale, FP8_DTYPE)
296 + at2 = auto_functionalized(
297 + MLA_ATTN_OP,
298 + q=q,
299 + kv_c_normed=kv_c_normed,
300 + k_pe=k_pe,
301 + output=output_attn,
302 + layer_name=layer_name,
303 + output_scale=input_scale,
304 + output_block_scale=output_scale_view,
305 + kv_cache_dummy_dep=kv_cache_dummy_dep,
306 + )
307 + return at2[1], at2[2]
308 +
309 + return _replacement_with_ln
310 +
311 def _replacement(
186 - q: torch.Tensor,
187 - kv_c_normed: torch.Tensor,
188 - k_pe: torch.Tensor,
189 - output_attn: torch.Tensor,
190 - input_scale: torch.Tensor,
191 - kv_cache_dummy_dep: torch.Tensor,
192 - ) -> tuple[torch.Tensor, torch.Tensor]:
312 + q, kv_c_normed, k_pe, output_attn, input_scale, kv_cache_dummy_dep
313 + ):
314 # MLA output in quant_dtype (FP4 packed as uint8)
315 output_attn = torch.empty(
316 [q.shape[0], self._output_dim // 2],
@@ -207,7 +328,7 @@ class MLAAttnNvfp4QuantPattern(
328 kv_c_normed=kv_c_normed,
329 k_pe=k_pe,
330 output=output_attn,
210 - layer_name=self._layer_name,
331 + layer_name=_ln,
332 output_scale=input_scale,
333 output_block_scale=output_scale_view,
334 kv_cache_dummy_dep=kv_cache_dummy_dep,
@@ -217,7 +338,7 @@ class MLAAttnNvfp4QuantPattern(
338 return _replacement
339
340 def get_inputs(self) -> list[torch.Tensor]:
220 - return [
341 + inputs: list = [
342 self.empty(5, self._num_heads, self._qk_head_dim, dtype=self._dtype),
343 self.empty(5, self._kv_lora_rank, dtype=self._dtype),
344 self.empty(5, 1, self._qk_rope_head_dim, dtype=self._dtype),
@@ -225,6 +346,9 @@ class MLAAttnNvfp4QuantPattern(
346 self.empty_fp32(1, 1),
347 self.empty(0, dtype=self._dtype),
348 ]
349 + if _USE_LAYERNAME:
350 + inputs.append(_encode_layer_name(self._layer_name))
351 + return inputs
352
353
354 class MLAAttnQuantFusionPass(VllmFusionPatternMatcherPass):
@@ -250,13 +374,19 @@ class MLAAttnQuantFusionPass(VllmFusionPatternMatcherPass):
374 "so no fusion patterns were registered."
375 )
376
377 + # When _USE_LAYERNAME is enabled, layer_name is a wildcard so all
378 + # layers produce the same pattern — register once then break.
379 for layer in layers:
380 if layer.impl.fused_output_quant_supported(kFp8StaticTensorSym):
381 self.register(MLAAttnFp8StaticQuantPattern(layer, dtype))
382 + if _USE_LAYERNAME:
383 + break
384
385 if current_platform.is_cuda() and hasattr(torch.ops._C, "scaled_fp4_quant"):
386 for layer in layers:
387 if layer.impl.fused_output_quant_supported(kNvfp4Dynamic):
388 self.register(MLAAttnNvfp4QuantPattern(layer, dtype))
389 + if _USE_LAYERNAME:
390 + break
391
392 self.dump_patterns(config, self.pm_pass)
vllm/compilation/passes/fusion/rope_kvcache_fusion.py
+72 -25
@@ -15,7 +15,13 @@ from vllm.model_executor.layers.attention.attention import (
15 Attention,
16 get_attention_context,
17 )
18 -from vllm.utils.torch_utils import direct_register_custom_op
18 +from vllm.utils.torch_utils import (
19 + _USE_LAYERNAME,
20 + LayerNameType,
21 + _encode_layer_name,
22 + _resolve_layer_name,
23 + direct_register_custom_op,
24 +)
25
26 from ..inductor_pass import enable_fake_mode
27 from ..vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
@@ -37,7 +43,7 @@ def fused_rope_and_unified_kv_cache_update_impl(
43 positions: torch.Tensor,
44 cos_sin_cache: torch.Tensor,
45 is_neox: bool,
40 - layer_name: str = "",
46 + layer_name: LayerNameType,
47 ) -> torch.Tensor:
48 """
49 This impl fetches the KV cache and slot mapping from the forward context,
@@ -46,6 +52,7 @@ def fused_rope_and_unified_kv_cache_update_impl(
52 that is passed to unified_attention to signal a side effect and
53 the data dependency between them to ensure torch.compile preserves ordering.
54 """
55 + layer_name = _resolve_layer_name(layer_name)
56 _, attn_layer, kv_cache, layer_slot_mapping = get_attention_context(layer_name)
57 if layer_slot_mapping is not None:
58 attn_layer.impl.do_rope_and_kv_cache_update(
@@ -70,7 +77,7 @@ def fused_rope_and_unified_kv_cache_update_fake(
77 positions: torch.Tensor,
78 cos_sin_cache: torch.Tensor,
79 is_neox: bool,
73 - layer_name: str = "",
80 + layer_name: LayerNameType,
81 ) -> torch.Tensor:
82 return torch.empty(0, device=query.device, dtype=query.dtype)
83
@@ -120,38 +127,30 @@ class RopeReshapeKVCachePattern:
127 num_kv_heads=self.num_kv_heads,
128 )
129
123 - def get_inputs(self) -> list[torch.Tensor]:
130 + def get_inputs(self) -> list:
131 # Sample inputs to help pattern tracing
132 T = 5
133 L = 4096
134 qkv = empty_bf16(T, self.q_size + self.k_size + self.v_size)
135 positions = empty_i64(T)
136 cos_sin_cache = empty_bf16(L, self.head_size)
130 - return [
131 - qkv,
132 - positions,
133 - cos_sin_cache,
134 - ]
137 + inputs: list = [qkv, positions, cos_sin_cache]
138 + if _USE_LAYERNAME:
139 + inputs.append(_encode_layer_name(self.layer_name))
140 + return inputs
141
136 - def register(self, pm_pass: PatternMatcherPass) -> None:
137 - def pattern(
138 - qkv: torch.Tensor,
139 - positions: torch.Tensor,
140 - cos_sin_cache: torch.Tensor,
141 - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
142 + def _mk_pattern_with_layer_name_input(self, _ln):
143 + """Pattern/replacement with layer_name as an explicit input."""
144 +
145 + def pattern(qkv, positions, cos_sin_cache, layer_name):
146 q, k, v = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1)
147 q, k = self.rope_matcher(positions, q, k, cos_sin_cache)
148 q = q.view(-1, self.num_heads, self.head_size)
149 k = k.view(-1, self.num_kv_heads, self.head_size)
150 v = v.view(-1, self.num_kv_heads, self.head_size_v)
147 - dummy = torch.ops.vllm.unified_kv_cache_update(k, v, self.layer_name)
148 - return dummy, q, k, v
149 -
150 - def replacement(
151 - qkv: torch.Tensor,
152 - positions: torch.Tensor,
153 - cos_sin_cache: torch.Tensor,
154 - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
151 + return torch.ops.vllm.unified_kv_cache_update(k, v, layer_name), q, k, v
152 +
153 + def replacement(qkv, positions, cos_sin_cache, layer_name):
154 q, k, v = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1)
155 q = q.view(-1, self.num_heads, self.head_size)
156 k = k.view(-1, self.num_kv_heads, self.head_size)
@@ -164,10 +163,50 @@ class RopeReshapeKVCachePattern:
163 positions=positions,
164 cos_sin_cache=cos_sin_cache,
165 is_neox=self.is_neox,
167 - layer_name=self.layer_name,
166 + layer_name=layer_name,
167 )
168 return results[0], results[1], results[2], v
169
170 + return pattern, replacement
171 +
172 + def _mk_pattern_with_layer_name_closure(self, _ln):
173 + """Pattern/replacement with layer_name as a closure constant."""
174 +
175 + def pattern(qkv, positions, cos_sin_cache):
176 + q, k, v = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1)
177 + q, k = self.rope_matcher(positions, q, k, cos_sin_cache)
178 + q = q.view(-1, self.num_heads, self.head_size)
179 + k = k.view(-1, self.num_kv_heads, self.head_size)
180 + v = v.view(-1, self.num_kv_heads, self.head_size_v)
181 + return torch.ops.vllm.unified_kv_cache_update(k, v, _ln), q, k, v
182 +
183 + def replacement(qkv, positions, cos_sin_cache):
184 + q, k, v = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1)
185 + q = q.view(-1, self.num_heads, self.head_size)
186 + k = k.view(-1, self.num_kv_heads, self.head_size)
187 + v = v.view(-1, self.num_kv_heads, self.head_size_v)
188 + results = auto_functionalized(
189 + self.FUSED_OP,
190 + query=q,
191 + key=k,
192 + value=v,
193 + positions=positions,
194 + cos_sin_cache=cos_sin_cache,
195 + is_neox=self.is_neox,
196 + layer_name=_ln,
197 + )
198 + return results[0], results[1], results[2], v
199 +
200 + return pattern, replacement
201 +
202 + def register(self, pm_pass: PatternMatcherPass) -> None:
203 + _ln = _encode_layer_name(self.layer_name)
204 +
205 + if _USE_LAYERNAME:
206 + pattern, replacement = self._mk_pattern_with_layer_name_input(_ln)
207 + else:
208 + pattern, replacement = self._mk_pattern_with_layer_name_closure(_ln)
209 +
210 # NOTE: use view_to_reshape to unify view/reshape to simplify
211 # pattern and increase matching opportunities
212 def fwd_and_view_to_reshape(*args, **kwargs) -> fx.GraphModule:
@@ -176,7 +215,11 @@ class RopeReshapeKVCachePattern:
215 return gm
216
217 pm.register_replacement(
179 - pattern, replacement, self.get_inputs(), fwd_and_view_to_reshape, pm_pass
218 + pattern,
219 + replacement,
220 + self.get_inputs(),
221 + fwd_and_view_to_reshape,
222 + pm_pass,
223 )
224
225
@@ -205,6 +248,8 @@ class RopeKVCacheFusionPass(VllmPatternMatcherPass):
248 self.max_token_num = cc.pass_config.rope_kvcache_fusion_max_token_num
249
250 attn_layers = get_layers_from_vllm_config(config, Attention)
251 + # When _USE_LAYERNAME is enabled, layer_name is a wildcard so all
252 + # layers produce the same pattern — register once then break.
253 for _, layer in attn_layers.items():
254 if layer.impl.fused_rope_kvcache_supported():
255 for is_neox in [True, False]:
@@ -212,6 +257,8 @@ class RopeKVCacheFusionPass(VllmPatternMatcherPass):
257 layer=layer,
258 is_neox=is_neox,
259 ).register(self.patterns)
260 + if _USE_LAYERNAME:
261 + break
262
263 self.dump_patterns(config, self.patterns)
264
vllm/envs.py
+4
@@ -129,6 +129,7 @@ if TYPE_CHECKING:
129 VLLM_ENABLE_V1_MULTIPROCESSING: bool = True
130 VLLM_LOG_BATCHSIZE_INTERVAL: float = -1
131 VLLM_DISABLE_COMPILE_CACHE: bool = False
132 + VLLM_USE_LAYERNAME: bool = True
133 Q_SCALE_CONSTANT: int = 200
134 K_SCALE_CONSTANT: int = 200
135 V_SCALE_CONSTANT: int = 100
@@ -1090,6 +1091,9 @@ environment_variables: dict[str, Callable[[], Any]] = {
1091 os.getenv("VLLM_LOG_BATCHSIZE_INTERVAL", "-1")
1092 ),
1093 "VLLM_DISABLE_COMPILE_CACHE": disable_compile_cache,
1094 + # If set to "0", disable LayerName opaque type for layer_name
1095 + # parameters in custom ops. Defaults to enabled on torch >= 2.11.
1096 + "VLLM_USE_LAYERNAME": lambda: bool(int(os.getenv("VLLM_USE_LAYERNAME", "1"))),
1097 # If set, vllm will run in development mode, which will enable
1098 # some additional endpoints for developing and debugging,
1099 # e.g. `/reset_prefix_cache`
vllm/model_executor/layers/attention/attention.py
+18 -9
@@ -25,6 +25,9 @@ from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
25 from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
26 from vllm.platforms import current_platform
27 from vllm.utils.torch_utils import (
28 + LayerNameType,
29 + _encode_layer_name,
30 + _resolve_layer_name,
31 direct_register_custom_op,
32 kv_cache_dtype_str_to_dtype,
33 )
@@ -414,7 +417,9 @@ class Attention(nn.Module, AttentionLayerBase):
417 `vllm.forward_context.get_forward_context().attn_metadata`.
418 """
419 if self.calculate_kv_scales:
417 - torch.ops.vllm.maybe_calc_kv_scales(query, key, value, self.layer_name)
420 + torch.ops.vllm.maybe_calc_kv_scales(
421 + query, key, value, _encode_layer_name(self.layer_name)
422 + )
423 output_dtype = query.dtype
424 if self.query_quant is not None:
425 # quantizing with a simple torch operation enables
@@ -466,6 +471,7 @@ class Attention(nn.Module, AttentionLayerBase):
471 )
472 else:
473 # Skip this if sharing KV cache with an earlier attention layer.
474 + encoded = _encode_layer_name(self.layer_name)
475 if (
476 not self.attn_backend.forward_includes_kv_cache_update
477 and self.kv_sharing_target_layer_name is None
@@ -473,14 +479,14 @@ class Attention(nn.Module, AttentionLayerBase):
479 and value is not None
480 ):
481 kv_cache_dummy_dep = torch.ops.vllm.unified_kv_cache_update(
476 - key, value, self.layer_name
482 + key, value, encoded
483 )
484 torch.ops.vllm.unified_attention_with_output(
485 query,
486 key,
487 value,
488 output,
483 - self.layer_name,
489 + encoded,
490 kv_cache_dummy_dep=kv_cache_dummy_dep,
491 )
492 return output.view(-1, hidden_size)
@@ -553,8 +559,9 @@ def maybe_calc_kv_scales(
559 query: torch.Tensor,
560 key: torch.Tensor,
561 value: torch.Tensor,
556 - layer_name: str,
562 + layer_name: LayerNameType,
563 ) -> None:
564 + layer_name = _resolve_layer_name(layer_name)
565 forward_context: ForwardContext = get_forward_context()
566 self = forward_context.no_compile_layers[layer_name]
567
@@ -570,7 +577,7 @@ def maybe_calc_kv_scales_fake(
577 query: torch.Tensor,
578 key: torch.Tensor,
579 value: torch.Tensor,
573 - layer_name: str,
580 + layer_name: LayerNameType,
581 ) -> None:
582 return
583
@@ -622,12 +629,13 @@ def get_attention_context(
629 def unified_kv_cache_update(
630 key: torch.Tensor,
631 value: torch.Tensor,
625 - layer_name: str,
632 + layer_name: LayerNameType,
633 ) -> torch.Tensor:
634 """
635 Returns a dummy that is passed to unified_attention to signal a side effect and
636 the data dependency between them to ensure torch.compile preserves ordering.
637 """
638 + layer_name = _resolve_layer_name(layer_name)
639 _, attn_layer, kv_cache, layer_slot_mapping = get_attention_context(layer_name)
640 if layer_slot_mapping is not None:
641 assert hasattr(attn_layer.impl, "do_kv_cache_update"), (
@@ -647,7 +655,7 @@ def unified_kv_cache_update(
655 def unified_kv_cache_update_fake(
656 key: torch.Tensor,
657 value: torch.Tensor,
650 - layer_name: str,
658 + layer_name: LayerNameType,
659 ) -> torch.Tensor:
660 return torch.empty(0, device=key.device, dtype=key.dtype)
661
@@ -666,7 +674,7 @@ def unified_attention_with_output(
674 key: torch.Tensor,
675 value: torch.Tensor,
676 output: torch.Tensor,
669 - layer_name: str,
677 + layer_name: LayerNameType,
678 output_scale: torch.Tensor | None = None,
679 output_block_scale: torch.Tensor | None = None,
680 kv_cache_dummy_dep: torch.Tensor | None = None,
@@ -675,6 +683,7 @@ def unified_attention_with_output(
683 # that ensures torch.compile preserves ordering between KV cache update and
684 # attention forward.
685 del kv_cache_dummy_dep
686 + layer_name = _resolve_layer_name(layer_name)
687 attn_metadata, self, kv_cache, _ = get_attention_context(layer_name)
688
689 self.impl.forward(
@@ -695,7 +704,7 @@ def unified_attention_with_output_fake(
704 key: torch.Tensor,
705 value: torch.Tensor,
706 output: torch.Tensor,
698 - layer_name: str,
707 + layer_name: LayerNameType,
708 output_scale: torch.Tensor | None = None,
709 output_block_scale: torch.Tensor | None = None,
710 kv_cache_dummy_dep: torch.Tensor | None = None,
vllm/model_executor/layers/attention/kv_transfer_utils.py
+2 -1
@@ -9,6 +9,7 @@ from vllm.distributed.kv_transfer import (
9 has_kv_transfer_group,
10 is_v1_kv_transfer_group,
11 )
12 +from vllm.utils.torch_utils import _resolve_layer_name
13
14
15 def maybe_transfer_kv_layer(func: Callable) -> Callable:
@@ -38,7 +39,7 @@ def maybe_transfer_kv_layer(func: Callable) -> Callable:
39 if not has_kv_transfer_group() or not is_v1_kv_transfer_group():
40 return func(*args, **kwargs)
41
41 - layer_name: str = args[layer_name_index]
42 + layer_name = _resolve_layer_name(args[layer_name_index])
43
44 # Extract attention context (metadata, layer, kv_cache, layer_slot_mapping)
45 attn_metadata, _, kv_cache, _ = get_attention_context(layer_name)
vllm/model_executor/layers/attention/mla_attention.py
+18 -7
@@ -240,6 +240,9 @@ from vllm.platforms import current_platform
240 from vllm.utils.flashinfer import has_flashinfer, has_nvidia_artifactory
241 from vllm.utils.math_utils import cdiv, round_down
242 from vllm.utils.torch_utils import (
243 + LayerNameType,
244 + _encode_layer_name,
245 + _resolve_layer_name,
246 direct_register_custom_op,
247 is_quantized_kv_cache,
248 kv_cache_dtype_str_to_dtype,
@@ -473,7 +476,12 @@ class MLAAttention(nn.Module, AttentionLayerBase):
476 output_shape: torch.Size | None = None,
477 ) -> torch.Tensor:
478 if self.calculate_kv_scales:
476 - torch.ops.vllm.maybe_calc_kv_scales(q, kv_c_normed, k_pe, self.layer_name)
479 + torch.ops.vllm.maybe_calc_kv_scales(
480 + q,
481 + kv_c_normed,
482 + k_pe,
483 + _encode_layer_name(self.layer_name),
484 + )
485
486 if self.use_direct_call:
487 forward_context: ForwardContext = get_forward_context()
@@ -505,10 +513,11 @@ class MLAAttention(nn.Module, AttentionLayerBase):
513 )
514 return output
515 else:
516 + encoded = _encode_layer_name(self.layer_name)
517 kv_cache_dummy_dep = torch.ops.vllm.unified_mla_kv_cache_update(
518 kv_c_normed,
519 k_pe,
511 - self.layer_name,
520 + encoded,
521 self.kv_cache_dtype,
522 self._k_scale,
523 )
@@ -518,7 +527,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
527 kv_c_normed,
528 k_pe,
529 output,
521 - self.layer_name,
530 + encoded,
531 kv_cache_dummy_dep=kv_cache_dummy_dep,
532 )
533 return output
@@ -900,7 +909,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
909 def unified_mla_kv_cache_update(
910 kv_c_normed: torch.Tensor,
911 k_pe: torch.Tensor,
903 - layer_name: str,
912 + layer_name: LayerNameType,
913 kv_cache_dtype: str,
914 k_scale: torch.Tensor,
915 ) -> torch.Tensor:
@@ -908,6 +917,7 @@ def unified_mla_kv_cache_update(
917 Returns a dummy that is passed to unified_attention to signal a side effect and
918 the data dependency between them to ensure torch.compile preserves ordering.
919 """
920 + layer_name = _resolve_layer_name(layer_name)
921 forward_context = get_forward_context()
922 attn_layer = forward_context.no_compile_layers[layer_name]
923 kv_cache = attn_layer.kv_cache
@@ -939,7 +949,7 @@ def unified_mla_kv_cache_update(
949 def unified_mla_kv_cache_update_fake(
950 kv_c_normed: torch.Tensor,
951 k_pe: torch.Tensor,
942 - layer_name: str,
952 + layer_name: LayerNameType,
953 kv_cache_dtype: str,
954 k_scale: torch.Tensor,
955 ) -> torch.Tensor:
@@ -959,7 +969,7 @@ def unified_mla_attention_with_output(
969 kv_c_normed: torch.Tensor,
970 k_pe: torch.Tensor,
971 output: torch.Tensor,
962 - layer_name: str,
972 + layer_name: LayerNameType,
973 output_scale: torch.Tensor | None = None,
974 output_block_scale: torch.Tensor | None = None,
975 kv_cache_dummy_dep: torch.Tensor | None = None,
@@ -968,6 +978,7 @@ def unified_mla_attention_with_output(
978 # that ensures torch.compile preserves ordering between KV cache update and
979 # attention forward.
980 del kv_cache_dummy_dep
981 + layer_name = _resolve_layer_name(layer_name)
982 attn_metadata, layer, kv_cache, _ = get_attention_context(layer_name)
983 layer.forward_impl(
984 q,
@@ -986,7 +997,7 @@ def unified_mla_attention_with_output_fake(
997 kv_c_normed: torch.Tensor,
998 k_pe: torch.Tensor,
999 output: torch.Tensor,
989 - layer_name: str,
1000 + layer_name: LayerNameType,
1001 output_scale: torch.Tensor | None = None,
1002 output_block_scale: torch.Tensor | None = None,
1003 kv_cache_dummy_dep: torch.Tensor | None = None,
vllm/model_executor/layers/attention/static_sink_attention.py
+12 -4
@@ -10,7 +10,12 @@ from vllm.logger import init_logger
10 from vllm.model_executor.custom_op import CustomOp
11 from vllm.model_executor.layers.attention import Attention
12 from vllm.utils.math_utils import cdiv
13 -from vllm.utils.torch_utils import direct_register_custom_op
13 +from vllm.utils.torch_utils import (
14 + LayerNameType,
15 + _encode_layer_name,
16 + _resolve_layer_name,
17 + direct_register_custom_op,
18 +)
19 from vllm.v1.attention.backend import (
20 AttentionBackend,
21 AttentionMetadata,
@@ -170,7 +175,9 @@ class StaticSinkAttention(Attention, CustomOp):
175 )
176 if not self.sink_populated:
177 self_kv_cache = self.kv_cache
173 - torch.ops.vllm.maybe_populate_sink(self_kv_cache, self.layer_name)
178 + torch.ops.vllm.maybe_populate_sink(
179 + self_kv_cache, _encode_layer_name(self.layer_name)
180 + )
181
182 return super().forward(query, key, value, output_shape)
183
@@ -224,8 +231,9 @@ class StaticSinkAttention(Attention, CustomOp):
231
232 def maybe_populate_sink(
233 self_kv_cache: torch.Tensor,
227 - layer_name: str,
234 + layer_name: LayerNameType,
235 ) -> None:
236 + layer_name = _resolve_layer_name(layer_name)
237 forward_context: ForwardContext = get_forward_context()
238 self = forward_context.no_compile_layers[layer_name]
239 if self.sink_populated or self_kv_cache.numel() == 0:
@@ -235,7 +243,7 @@ def maybe_populate_sink(
243
244 def maybe_populate_sink_fake(
245 self_kv_cache: torch.Tensor,
238 - layer_name: str,
246 + layer_name: LayerNameType,
247 ) -> None:
248 return
249
vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py
+11 -11
@@ -32,15 +32,15 @@ from vllm.model_executor.layers.fused_moe.runner.shared_experts import (
32 )
33 from vllm.platforms import current_platform
34 from vllm.utils.torch_utils import (
35 - HAS_OPAQUE_TYPE,
36 - ModuleName,
35 + _USE_LAYERNAME,
36 + LayerName,
37 direct_register_custom_op,
38 )
39
40
41 def get_layer_from_name(layer_name: str) -> torch.nn.Module:
42 forward_context: ForwardContext = get_forward_context()
43 - if not HAS_OPAQUE_TYPE and layer_name == "from_forward_context":
43 + if not _USE_LAYERNAME and layer_name == "from_forward_context":
44 all_moe_layers = forward_context.all_moe_layers
45 assert all_moe_layers is not None
46 moe_layer_index = forward_context.moe_layer_index
@@ -55,21 +55,21 @@ def get_layer_from_name(layer_name: str) -> torch.nn.Module:
55 return forward_context.no_compile_layers[layer_name]
56
57
58 -# On torch >= 2.11, layer_name is a hoisted ModuleName opaque object;
58 +# On torch >= 2.11, layer_name is a hoisted LayerName opaque object;
59 # on older versions it remains a plain str.
60 if TYPE_CHECKING:
61 from typing import TypeAlias
62
63 - _layer_name_type: TypeAlias = str | ModuleName
63 + _layer_name_type: TypeAlias = str | LayerName
64 else:
65 - _layer_name_type = ModuleName if HAS_OPAQUE_TYPE else str
65 + _layer_name_type = LayerName if _USE_LAYERNAME else str
66
67
68 @torch.compiler.assume_constant_result
69 -def _resolve_layer_name(layer_name: str | ModuleName) -> str:
69 +def _resolve_layer_name(layer_name: str | LayerName) -> str:
70 from torch._library.fake_class_registry import FakeScriptObject
71
72 - if isinstance(layer_name, ModuleName):
72 + if isinstance(layer_name, LayerName):
73 return layer_name.value
74 elif isinstance(layer_name, FakeScriptObject):
75 return layer_name.real_obj.value
@@ -331,9 +331,9 @@ class MoERunnerBase(MoERunner):
331 assert len(trunc_sizes) == 1
332 return func(states, trunc_sizes[0])
333
334 - def _encode_layer_name(self) -> str | ModuleName:
335 - if HAS_OPAQUE_TYPE:
336 - return ModuleName(self.layer_name)
334 + def _encode_layer_name(self) -> str | LayerName:
335 + if _USE_LAYERNAME:
336 + return LayerName(self.layer_name)
337 # Can be unavailable or None in unittests
338 if (
339 is_forward_context_available()
vllm/model_executor/layers/mamba/gdn_linear_attn.py
+10 -4
@@ -56,7 +56,12 @@ from vllm.model_executor.utils import set_weight_attrs
56 from vllm.platforms import current_platform
57 from vllm.transformers_utils.configs.qwen3_next import Qwen3NextConfig
58 from vllm.triton_utils import tl, triton
59 -from vllm.utils.torch_utils import direct_register_custom_op
59 +from vllm.utils.torch_utils import (
60 + LayerNameType,
61 + _encode_layer_name,
62 + _resolve_layer_name,
63 + direct_register_custom_op,
64 +)
65 from vllm.v1.attention.backend import AttentionMetadata
66 from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
67
@@ -568,7 +573,7 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
573 b,
574 a,
575 core_attn_out,
571 - self.prefix,
576 + _encode_layer_name(self.prefix),
577 )
578
579 # ============================================================
@@ -1084,13 +1089,14 @@ def gdn_attention_core(
1089 b: torch.Tensor,
1090 a: torch.Tensor,
1091 core_attn_out: torch.Tensor,
1087 - layer_name: str,
1092 + layer_name: LayerNameType,
1093 ) -> None:
1094 """
1095 Custom op for the core attention computation.
1096 Only handles the convolution + recurrent attention part.
1097 Input/output projections are handled outside this op.
1098 """
1099 + layer_name = _resolve_layer_name(layer_name)
1100 forward_context: ForwardContext = get_forward_context()
1101 self = forward_context.no_compile_layers[layer_name]
1102 self._forward_core(
@@ -1106,7 +1112,7 @@ def gdn_attention_core_fake(
1112 b: torch.Tensor,
1113 a: torch.Tensor,
1114 core_attn_out: torch.Tensor,
1109 - layer_name: str,
1115 + layer_name: LayerNameType,
1116 ) -> None:
1117 """Fake implementation for torch.compile."""
1118 return
vllm/model_executor/layers/mamba/mamba_mixer.py
+10 -4
@@ -36,7 +36,12 @@ from vllm.model_executor.layers.mamba.ops.mamba_ssm import (
36 )
37 from vllm.model_executor.utils import set_weight_attrs
38 from vllm.platforms import current_platform
39 -from vllm.utils.torch_utils import direct_register_custom_op
39 +from vllm.utils.torch_utils import (
40 + LayerNameType,
41 + _encode_layer_name,
42 + _resolve_layer_name,
43 + direct_register_custom_op,
44 +)
45 from vllm.v1.attention.backends.mamba1_attn import Mamba1AttentionMetadata
46
47
@@ -228,7 +233,7 @@ class MambaMixer(MambaBase, PluggableLayer):
233 torch.ops.vllm.mamba_mixer(
234 hidden_states,
235 output,
231 - self.prefix,
236 + _encode_layer_name(self.prefix),
237 )
238
239 def forward_impl(self, hidden_states: torch.Tensor, output: torch.Tensor):
@@ -515,8 +520,9 @@ def split_batch_to_prefill_and_decode(
520 def mamba_mixer(
521 hidden_states: torch.Tensor,
522 output: torch.Tensor,
518 - layer_name: str,
523 + layer_name: LayerNameType,
524 ) -> None:
525 + layer_name = _resolve_layer_name(layer_name)
526 forward_context: ForwardContext = get_forward_context()
527 self = forward_context.no_compile_layers[layer_name]
528 self.forward_impl(hidden_states=hidden_states, output=output)
@@ -525,7 +531,7 @@ def mamba_mixer(
531 def mamba_mixer_fake(
532 hidden_states: torch.Tensor,
533 output: torch.Tensor,
528 - layer_name: str,
534 + layer_name: LayerNameType,
535 ) -> None:
536 return
537
vllm/model_executor/layers/mamba/mamba_mixer2.py
+10 -4
@@ -44,7 +44,12 @@ from vllm.model_executor.model_loader.weight_utils import (
44 from vllm.model_executor.parameter import BasevLLMParameter
45 from vllm.model_executor.utils import set_weight_attrs
46 from vllm.platforms import current_platform
47 -from vllm.utils.torch_utils import direct_register_custom_op
47 +from vllm.utils.torch_utils import (
48 + LayerNameType,
49 + _encode_layer_name,
50 + _resolve_layer_name,
51 + direct_register_custom_op,
52 +)
53 from vllm.v1.attention.backend import AttentionMetadata
54 from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadata
55
@@ -536,7 +541,7 @@ class MambaMixer2(MambaBase, PluggableLayer):
541 torch.ops.vllm.mamba_mixer2(
542 projected_states,
543 ssm_output,
539 - self.prefix,
544 + _encode_layer_name(self.prefix),
545 )
546
547 # 4. gated MLP
@@ -927,8 +932,9 @@ class MambaMixer2(MambaBase, PluggableLayer):
932 def mamba_mixer2(
933 projected_states: torch.Tensor,
934 output: torch.Tensor,
930 - layer_name: str,
935 + layer_name: LayerNameType,
936 ) -> None:
937 + layer_name = _resolve_layer_name(layer_name)
938 forward_context: ForwardContext = get_forward_context()
939 self = forward_context.no_compile_layers[layer_name]
940 self.conv_ssm_forward(projected_states=projected_states, output=output)
@@ -937,7 +943,7 @@ def mamba_mixer2(
943 def mamba_mixer2_fake(
944 projected_states: torch.Tensor,
945 output: torch.Tensor,
940 - layer_name: str,
946 + layer_name: LayerNameType,
947 ) -> None:
948 return
949
vllm/model_executor/layers/sparse_attn_indexer.py
+11 -5
@@ -11,7 +11,12 @@ from vllm.logger import init_logger
11 from vllm.model_executor.custom_op import CustomOp
12 from vllm.platforms import current_platform
13 from vllm.utils.deep_gemm import fp8_mqa_logits, fp8_paged_mqa_logits, has_deep_gemm
14 -from vllm.utils.torch_utils import direct_register_custom_op
14 +from vllm.utils.torch_utils import (
15 + LayerNameType,
16 + _encode_layer_name,
17 + _resolve_layer_name,
18 + direct_register_custom_op,
19 +)
20 from vllm.v1.attention.backends.mla.indexer import (
21 DeepseekV32IndexerMetadata,
22 )
@@ -30,7 +35,7 @@ RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024
35
36 def sparse_attn_indexer(
37 hidden_states: torch.Tensor,
33 - k_cache_prefix: str,
38 + k_cache_prefix: LayerNameType,
39 kv_cache: torch.Tensor,
40 q_fp8: torch.Tensor,
41 k: torch.Tensor,
@@ -46,6 +51,7 @@ def sparse_attn_indexer(
51 # careful! this will be None in dummy run
52 attn_metadata = get_forward_context().attn_metadata
53 fp8_dtype = current_platform.fp8_dtype()
54 + k_cache_prefix = _resolve_layer_name(k_cache_prefix)
55
56 # assert isinstance(attn_metadata, dict)
57 if not isinstance(attn_metadata, dict):
@@ -253,7 +259,7 @@ def sparse_attn_indexer(
259
260 def sparse_attn_indexer_fake(
261 hidden_states: torch.Tensor,
256 - k_cache_prefix: str,
262 + k_cache_prefix: LayerNameType,
263 kv_cache: torch.Tensor,
264 q_fp8: torch.Tensor,
265 k: torch.Tensor,
@@ -342,7 +348,7 @@ class SparseAttnIndexer(CustomOp):
348 ):
349 return torch.ops.vllm.sparse_attn_indexer(
350 hidden_states,
345 - self.k_cache.prefix,
351 + _encode_layer_name(self.k_cache.prefix),
352 self.k_cache.kv_cache,
353 q_fp8,
354 k,
@@ -366,7 +372,7 @@ class SparseAttnIndexer(CustomOp):
372 if rocm_aiter_ops.is_enabled():
373 return torch.ops.vllm.rocm_aiter_sparse_attn_indexer(
374 hidden_states,
369 - self.k_cache.prefix,
375 + _encode_layer_name(self.k_cache.prefix),
376 self.k_cache.kv_cache,
377 q_fp8,
378 k,
vllm/utils/torch_utils.py
+31 -5
@@ -15,6 +15,7 @@ from packaging import version
15 from packaging.version import Version
16 from torch.library import Library, infer_schema
17
18 +import vllm.envs as envs
19 from vllm.logger import init_logger
20
21 if TYPE_CHECKING:
@@ -709,37 +710,62 @@ def is_torch_equal(target: str) -> bool:
710
711 HAS_OPAQUE_TYPE = is_torch_equal_or_newer("2.11.0.dev")
712
713 +# Allow toggling LayerName usage via environment variable.
714 +# Defaults to True on torch >= 2.11, False otherwise.
715 +# Set VLLM_USE_LAYERNAME=0 to disable even on torch >= 2.11.
716 +_USE_LAYERNAME = HAS_OPAQUE_TYPE and envs.VLLM_USE_LAYERNAME
717 +
718 if HAS_OPAQUE_TYPE:
719 from torch._opaque_base import OpaqueBase
720 else:
721 OpaqueBase = object # type: ignore[misc, assignment]
722
723
718 -class ModuleName(OpaqueBase): # type: ignore[misc]
724 +class LayerName(OpaqueBase): # type: ignore[misc]
725 """Wraps a module name string for use as a torch opaque type.
726
727 When torch >= 2.11, this is registered as a hoisted value-type opaque
728 object so that torch.compile lifts it as a graph input instead of baking
723 - it as a constant. This avoids per-layer recompilation for MOE ops.
729 + it as a constant. This avoids per-layer recompilation for custom ops
730 + that accept layer name strings (attention, MOE, KV cache, etc.).
731 """
732
733 def __init__(self, value: str):
734 self.value = value
735
736 def __eq__(self, other):
730 - return isinstance(other, ModuleName) and self.value == other.value
737 + return isinstance(other, LayerName) and self.value == other.value
738
739 def __hash__(self):
740 return hash(self.value)
741
742 def __fx_repr__(self):
736 - return (f"ModuleName({self.value!r})", {ModuleName})
743 + return (f"LayerName({self.value!r})", {"LayerName": LayerName})
744
745
746 if HAS_OPAQUE_TYPE:
747 from torch._library.opaque_object import register_opaque_type
748
742 - register_opaque_type(ModuleName, typ="value", hoist=True)
749 + register_opaque_type(LayerName, typ="value", hoist=True)
750 +
751 +# On torch >= 2.11 (with VLLM_USE_LAYERNAME enabled), custom op
752 +# layer_name parameters use LayerName; otherwise they remain plain str.
753 +if TYPE_CHECKING:
754 + from typing import TypeAlias
755 +
756 + LayerNameType: TypeAlias = str | LayerName
757 +else:
758 + LayerNameType = LayerName if _USE_LAYERNAME else str
759 +
760 +
761 +def _resolve_layer_name(layer_name: str | LayerName) -> str:
762 + """Unwrap a LayerName to str, or return str unchanged."""
763 + return layer_name.value if isinstance(layer_name, LayerName) else layer_name
764 +
765 +
766 +def _encode_layer_name(layer_name: str) -> str | LayerName:
767 + """Wrap a str layer name as LayerName when enabled."""
768 + return LayerName(layer_name) if _USE_LAYERNAME else layer_name
769
770
771 # Supports xccl with PyTorch versions >= 2.8.0.dev for XPU platform
vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
+6 -2
@@ -9,6 +9,7 @@ import torch
9 from vllm.forward_context import get_forward_context
10 from vllm.platforms import current_platform
11 from vllm.triton_utils import tl, triton
12 +from vllm.utils.torch_utils import LayerNameType
13 from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerMetadata
14 from vllm.v1.attention.ops.common import pack_seq_triton, unpack_seq_triton
15
@@ -459,7 +460,7 @@ def rocm_fp8_mqa_logits(
460
461 def rocm_aiter_sparse_attn_indexer_fake(
462 hidden_states: torch.Tensor,
462 - k_cache_prefix: str,
463 + k_cache_prefix: LayerNameType,
464 kv_cache: torch.Tensor,
465 q_fp8: torch.Tensor,
466 k: torch.Tensor,
@@ -486,7 +487,7 @@ def rocm_aiter_sparse_attn_indexer_fake(
487
488 def rocm_aiter_sparse_attn_indexer(
489 hidden_states: torch.Tensor,
489 - k_cache_prefix: str,
490 + k_cache_prefix: LayerNameType,
491 kv_cache: torch.Tensor,
492 q_fp8: torch.Tensor,
493 k: torch.Tensor,
@@ -502,6 +503,9 @@ def rocm_aiter_sparse_attn_indexer(
503 # careful! this will be None in dummy run
504 attn_metadata = get_forward_context().attn_metadata
505 fp8_dtype = current_platform.fp8_dtype()
506 + from vllm.utils.torch_utils import _resolve_layer_name
507 +
508 + k_cache_prefix = _resolve_layer_name(k_cache_prefix)
509 # assert isinstance(attn_metadata, dict)
510 if not isinstance(attn_metadata, dict):
511 return rocm_aiter_sparse_attn_indexer_fake(