@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 92e8518d37

Improve Transformers modelling backend `fx` tracer (#49957)

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

Harry Mellor committed Jul 27, 2026 at 12:51 UTC 92e8518d376c3c9a6979da4d56ee6263ef6af630

9 files changed +307 -128

tests/models/transformers/fusers/test_linear.py

+3 -4

@@ -368,12 +368,12 @@ def test_qkv_identifies_output_projection():

368 assert get_fuser(PerHeadQKNormAttention()).o_name == "o_proj"

369

370

371 -def test_fuser_is_cached_per_class():

371 +def test_fuser_is_cached_per_class_and_structure():

372 with torch.device("meta"):

373 fuser_a = get_fuser(GLUMLP())

374 fuser_b = get_fuser(GLUMLP())

375 assert fuser_a is fuser_b

376 - assert GLUMLP in get_fuser.cache

376 + assert any(key[0] is GLUMLP for key in get_fuser.cache)

377

378

379 @pytest.mark.parametrize("cls", [NotAnMLP, UntraceableMLP])

@@ -460,8 +460,7 @@ def test_unfusable_modules_are_not_fused(cls, default_vllm_config):

460 fuser = get_fuser(module)

461 # Either no pattern matches the class, or this instance fails validation

462 # (`recursive_replace` gates fusion and its weight mappings on `validate`)

463 - model_config = default_vllm_config.model_config

464 - assert fuser is None or not fuser.validate(module, model_config)

463 + assert fuser is None or not fuser.validate(module, default_vllm_config)

464

465

466 def test_act_and_mul_derived_from_module(default_vllm_config):

tests/models/transformers/fusers/test_rms_norm.py

+11 -6

@@ -148,11 +148,13 @@ def test_rms_norm_builds_vllm_class(cls, expected, zero_centered, default_vllm_c

148

149 # `default_vllm_config` supplies the config context the CustomOp needs; the

150 # weightless path reads hidden size from the model config, so stub it.

151 - model_config = SimpleNamespace(get_hidden_size=lambda: 16)

151 + vllm_config = SimpleNamespace(

152 + model_config=SimpleNamespace(get_hidden_size=lambda: 16)

153 + )

154 with torch.device("meta"):

155 module = cls()

156 fuser = get_fuser(module)

155 - built = fuser.fuse(module, "norm", model_config, None)

157 + built = fuser.fuse(module, "norm", vllm_config)

158 from vllm.model_executor.models.transformers.fusers.rms_norm import (

159 TPAwareNormMixin,

160 )

@@ -176,8 +178,9 @@ def test_fused_rms_norm_op_default_eps(default_vllm_config):

178 fuser = get_fuser(module)

179 assert isinstance(fuser, RMSNormFuser)

180 assert not fuser.zero_centered

179 - model_config = SimpleNamespace(get_hidden_size=lambda: 16, dtype=torch.float32)

180 - built = fuser.fuse(module, "norm", model_config, None)

181 + mc = SimpleNamespace(get_hidden_size=lambda: 16, dtype=torch.float32)

182 + vllm_config = SimpleNamespace(model_config=mc)

183 + built = fuser.fuse(module, "norm", vllm_config)

184 assert isinstance(built, VLLMRMSNorm)

185 assert built.variance_epsilon == torch.finfo(torch.float32).eps

186

@@ -185,11 +188,13 @@ def test_fused_rms_norm_op_default_eps(default_vllm_config):

188 def test_eps_is_derived_per_instance(default_vllm_config):

189 """Two instances of the same norm class with different eps must fuse to their

190 own eps: the type-cached fuser holds only structure, not this value."""

188 - model_config = SimpleNamespace(get_hidden_size=lambda: 16)

191 + vllm_config = SimpleNamespace(

192 + model_config=SimpleNamespace(get_hidden_size=lambda: 16)

193 + )

194 with torch.device("meta"):

195 for eps in (1e-5, 1e-6):

196 module = RMSNorm(16, eps=eps)

192 - built = get_fuser(module).fuse(module, "norm", model_config, None)

197 + built = get_fuser(module).fuse(module, "norm", vllm_config)

198 assert built.variance_epsilon == eps

199

200

vllm/model_executor/models/transformers/base.py

+4 -5

@@ -104,6 +104,7 @@ class Base(

104 super().__init__()

105 logger.info("Using Transformers modeling backend.")

106

107 + self.vllm_config = vllm_config

108 self.config = vllm_config.model_config.hf_config

109 self.text_config = self.config.get_text_config()

110 self.cache_config = vllm_config.cache_config

@@ -357,7 +358,7 @@ class Base(

358 tip = get_feature_request_tip(

359 self.model_config.model, self.model_config.trust_remote_code

360 )

360 - logger.warning(

361 + logger.warning_once(

362 "%s does not define a pipeline parallel plan. The Transformers "

363 "modeling backend will infer the split from the layers of %s in order "

364 "of declaration and keep parameter-free modules on every rank. This "

@@ -452,7 +453,7 @@ class Base(

453 # Prefix the patterns because we always start from `self.model`

454 tp_plan = {maybe_prefix("model", k): v for k, v in tp_plan.items()}

455 # Detect fusable patterns once per module class (cached, so this is cheap)

455 - fusers = Fusers(self.model, self.model_config)

456 + fusers = Fusers(self.model, self.vllm_config)

457

458 def register_fusion(fuser: BaseFuser, prefix: str):

459 """Register a fused layer's mappings just before it is built."""

@@ -503,9 +504,7 @@ class Base(

504 new_module = replace_conv_class(child_module)

505 elif (fuser := fusers[child_module]) is not None:

506 register_fusion(fuser, qual_name)

506 - new_module = fuser.fuse(

507 - child_module, qual_name, self.model_config, self.quant_config

508 - )

507 + new_module = fuser.fuse(child_module, qual_name, self.vllm_config)

508 logger.info_once(fuser.info(child_name))

509 _recursive_replace(new_module, prefix=qual_name)

510 elif not isinstance(child_module, MoERunner):

vllm/model_executor/models/transformers/fuser.py

+17 -9

@@ -25,14 +25,19 @@ from vllm.model_executor.models.transformers.fusers import (

25 from vllm.model_executor.models.transformers.fx_utils import trace

27 if TYPE_CHECKING:

28 - from vllm.config.model import ModelConfig

28 + from vllm.config import VllmConfig

30 logger = init_logger(__name__)

33 -@cached(cache={}, key=type)

33 +def key(module: nn.Module) -> tuple:

34 + """Cache key for `get_fuser`. Considers module type and its immediate children."""

35 + return (type(module), tuple(name for name, _ in module.named_children()))

36 +

37 +

38 +@cached(cache={}, key=key)

39 def get_fuser(module: nn.Module) -> BaseFuser | None:

35 - """The fuser for `type(module)` (cached per class), or `None` if no match."""

40 + """The fuser for `module`'s class and shape (cached), or `None` if no match."""

41 # Projection fusions need >=2 sibling linears; the RMSNorm fusion needs a

42 # leaf module (raw tensor math, no submodules). Nothing else can match, and

43 # tracing is skipped for it.

@@ -48,11 +53,14 @@ def get_fuser(module: nn.Module) -> BaseFuser | None:

53 try:

54 fuser.update_forward(module)

55 except Exception as exc:

51 - # An unrecognised source just means we cannot fuse here.

56 logger.debug(

53 - "Could not rewrite %s for fusion: %s", type(module), exc

57 + "Attempted to fuse %s using %s but failed "

58 + "to update its forward method: %s",

59 + type(module),

60 + fuser_cls.__name__,

61 + exc,

62 )

55 - return None

63 + continue

64 return fuser

65 # A norm we could not match structurally is left unfused; flag likely misses.

66 if module.__class__.__name__.endswith("RMSNorm"):

@@ -67,12 +75,12 @@ def get_fuser(module: nn.Module) -> BaseFuser | None:

75 class Fusers(UserDict):

76 """Mapping from module class to fuser, for all fusable classes in a model."""

70 - def __init__(self, model: nn.Module, model_config: "ModelConfig"):

71 - self.model_config = model_config

78 + def __init__(self, model: nn.Module, vllm_config: "VllmConfig"):

79 + self.vllm_config = vllm_config

80 super().__init__({type(m): get_fuser(m) for m in model.modules()})

82 def __getitem__(self, m: nn.Module) -> BaseFuser | None:

83 fuser = self.data.get(type(m))

76 - if fuser is not None and fuser.validate(m, self.model_config):

84 + if fuser is not None and fuser.validate(m, self.vllm_config):

85 return fuser

86 return None

vllm/model_executor/models/transformers/fusers/base.py

+6 -19

@@ -13,8 +13,7 @@ from torch import fx, nn

13 from vllm.model_executor.models.utils import ShardId, maybe_prefix

15 if TYPE_CHECKING:

16 - from vllm.config.model import ModelConfig

17 - from vllm.model_executor.layers.quantization import QuantizationConfig

16 + from vllm.config import VllmConfig

19 @dataclass

@@ -36,16 +35,12 @@ class BaseFuser(ABC):

35 """Match the pattern in `graph`, returning a fuser if found."""

37 @abstractmethod

39 - def validate(self, module: nn.Module, model_config: "ModelConfig") -> bool:

38 + def validate(self, module: nn.Module, vllm_config: "VllmConfig") -> bool:

39 """Whether this fuser can be applied to this `module` instance."""

41 @abstractmethod

42 def fuse(

44 - self,

45 - module: nn.Module,

46 - prefix: str,

47 - model_config: "ModelConfig",

48 - quant_config: "QuantizationConfig",

43 + self, module: nn.Module, prefix: str, vllm_config: "VllmConfig"

44 ) -> nn.Module:

45 """Apply the fusion to an already-validated `module`, returning the

46 module to install in its place (mutated in place, or freshly built)."""

@@ -123,24 +118,16 @@ class StackedFuser(BaseFuser):

118

119 @abstractmethod

120 def update_attrs(

126 - self,

127 - module: nn.Module,

128 - prefix: str,

129 - model_config: "ModelConfig",

130 - quant_config: "QuantizationConfig",

121 + self, module: nn.Module, prefix: str, vllm_config: "VllmConfig"

122 ) -> None:

123 """Replace `module`'s submodules with the merged module."""

124

125 def fuse(

135 - self,

136 - module: nn.Module,

137 - prefix: str,

138 - model_config: "ModelConfig",

139 - quant_config: "QuantizationConfig",

126 + self, module: nn.Module, prefix: str, vllm_config: "VllmConfig"

127 ) -> nn.Module:

128 """Fuse an already-validated `module` in place (see `Fusers.__getitem__`).

129

130 Builds the merged submodule and binds the compiled forward."""

144 - self.update_attrs(module, prefix, model_config, quant_config)

131 + self.update_attrs(module, prefix, vllm_config)

132 module.forward = types.MethodType(self.fused_forward, module)

133 return module

vllm/model_executor/models/transformers/fusers/glu.py

+4 -8

@@ -33,8 +33,7 @@ from vllm.model_executor.models.transformers.utils import (

33 from vllm.model_executor.models.utils import ShardId, maybe_prefix

35 if TYPE_CHECKING:

36 - from vllm.config.model import ModelConfig

37 - from vllm.model_executor.layers.quantization import QuantizationConfig

36 + from vllm.config import VllmConfig

38 logger = init_logger(__name__)

@@ -168,7 +167,7 @@ class GLUFuser(StackedFuser):

167 replace_expr(funcdef, muls[0], act_call)

168 self.fused_forward = compile_forward(funcdef, fn)

169

171 - def validate(self, module: nn.Module, model_config: "ModelConfig") -> bool:

170 + def validate(self, module: nn.Module, vllm_config: "VllmConfig") -> bool:

171 act = module.get_submodule(self.act_name)

172 if self._get_act_and_mul_name(act) is None:

173 logger.debug("No AndMul equivalent for %s; skipping fusion", type(act))

@@ -176,12 +175,9 @@ class GLUFuser(StackedFuser):

175 return True

176

177 def update_attrs(

179 - self,

180 - module: nn.Module,

181 - prefix: str,

182 - model_config: "ModelConfig",

183 - quant_config: "QuantizationConfig",

178 + self, module: nn.Module, prefix: str, vllm_config: "VllmConfig"

179 ) -> None:

180 + quant_config = vllm_config.quant_config

181 act_fn = self._get_act_and_mul(module.get_submodule(self.act_name))

182 gate = module.get_submodule(self.gate_name)

183 up = module.get_submodule(self.up_name)

vllm/model_executor/models/transformers/fusers/qkv.py

+17 -19

@@ -17,6 +17,7 @@ from vllm.model_executor.models.transformers.fx_utils import (

17 is_linear,

18 recover_forward,

19 replace_expr,

20 + returned_linear,

21 single_self_call,

22 )

23 from vllm.model_executor.models.transformers.utils import (

@@ -26,8 +27,7 @@ from vllm.model_executor.models.transformers.utils import (

27 from vllm.model_executor.models.utils import ShardId, maybe_prefix

29 if TYPE_CHECKING:

29 - from vllm.config.model import ModelConfig

30 - from vllm.model_executor.layers.quantization import QuantizationConfig

30 + from vllm.config import VllmConfig

32 logger = init_logger(__name__)

@@ -84,15 +84,16 @@ class QKVFuser(StackedFuser):

84 return None

85 q, k, v = qkv_nodes

86 names = dict(q_name=q.target, k_name=k.target, v_name=v.target)

87 - attn_width = module.get_submodule(q.target).out_features

88 - candidates = [

89 - name

90 - for name, child in module.named_children()

91 - if isinstance(child, nn.Linear)

92 - and name not in names.values()

93 - and child.in_features == attn_width

94 - ]

95 - names["o_name"] = candidates[0] if len(candidates) == 1 else None

87 + # o_proj produces the module's output.

88 + o_name = returned_linear(graph, module)

89 + # o_proj must be compatible with the q/k/v projections.

90 + if o_name in names.values() or (

91 + o_name is not None

92 + and module.get_submodule(o_name).in_features

93 + != module.get_submodule(q.target).out_features

94 + ):

95 + o_name = None

96 + names["o_name"] = o_name

97 return cls(source_cls=type(module).__name__, **names)

99 def update_forward(self, module: nn.Module) -> None:

@@ -149,12 +150,12 @@ class QKVFuser(StackedFuser):

150 replace_expr(funcdef, call, ast.Name(id=temp, ctx=ast.Load()))

151 self.fused_forward = compile_forward(funcdef, fn)

152

152 - def validate(self, module: nn.Module, model_config: "ModelConfig") -> bool:

153 + def validate(self, module: nn.Module, vllm_config: "VllmConfig") -> bool:

154 """Shapes must be compatible for a single merged, head-sharded GEMM."""

155 q = module.get_submodule(self.q_name)

156 k = module.get_submodule(self.k_name)

157 v = module.get_submodule(self.v_name)

157 - head_size = model_config.get_head_size()

158 + head_size = vllm_config.model_config.get_head_size()

159 compatible = (

160 q.in_features == k.in_features == v.in_features

161 and len({proj.bias is None for proj in (q, k, v)}) == 1

@@ -167,13 +168,10 @@ class QKVFuser(StackedFuser):

168 return compatible

169

170 def update_attrs(

170 - self,

171 - module: nn.Module,

172 - prefix: str,

173 - model_config: "ModelConfig",

174 - quant_config: "QuantizationConfig",

171 + self, module: nn.Module, prefix: str, vllm_config: "VllmConfig"

172 ) -> None:

176 - head_size = model_config.get_head_size()

173 + quant_config = vllm_config.quant_config

174 + head_size = vllm_config.model_config.get_head_size()

175 q = module.get_submodule(self.q_name)

176 k = module.get_submodule(self.k_name)

177 merged = QKVParallelLinear(

vllm/model_executor/models/transformers/fusers/rms_norm.py

+7 -12

@@ -21,13 +21,13 @@ from vllm.model_executor.models.transformers.fx_utils import (

21 find_node,

22 forward_input_count,

23 is_op,

24 + output_value,

25 peel,

26 trace,

27 )

29 if TYPE_CHECKING:

29 - from vllm.config.model import ModelConfig

30 - from vllm.model_executor.layers.quantization import QuantizationConfig

30 + from vllm.config import VllmConfig

33 def _is_squared(node: object, x: fx.Node) -> bool:

@@ -71,10 +71,8 @@ def _is_one_plus(node: object) -> bool:

72 def _has_trailing_compute(graph: fx.Graph, node: fx.Node) -> bool:

73 """Does the forward compute anything after `node` before returning?"""

74 - output = find_node(graph, lambda n: n.op == "output")

75 - if output is None or not output.args:

76 - return False

77 - return peel(output.args[0]) is not node

74 + value = output_value(graph)

75 + return value is not None and peel(value) is not node

78 class TPAwareNormMixin(nn.Module):

@@ -185,17 +183,14 @@ class RMSNormFuser(BaseFuser):

183 return eps

184 return None

185

188 - def validate(self, module: nn.Module, model_config: "ModelConfig") -> bool:

186 + def validate(self, module: nn.Module, vllm_config: "VllmConfig") -> bool:

187 return True

188

189 def fuse(

192 - self,

193 - module: nn.Module,

194 - prefix: str,

195 - model_config: "ModelConfig",

196 - quant_config: "QuantizationConfig",

190 + self, module: nn.Module, prefix: str, vllm_config: "VllmConfig"

191 ) -> nn.Module:

192 """Fuse the matched RMSNorm pattern into a vLLM fused RMSNorm CustomOp."""

193 + model_config = vllm_config.model_config

194 weight = getattr(module, "weight", None)

195 hidden_size = (

196 weight.size(0) if weight is not None else model_config.get_hidden_size()

vllm/model_executor/models/transformers/fx_utils.py

+238 -46

@@ -9,10 +9,13 @@ stays live Python. `fusion.py` builds the concrete fusion patterns on top.

9 """

11 import ast

12 +import contextlib

13 import inspect

14 import operator

15 import textwrap

16 from collections.abc import Callable

17 +from itertools import chain

18 +from unittest import mock

20 import torch

21 from torch import fx, nn

@@ -22,46 +25,64 @@ from vllm.logger import init_logger

26 logger = init_logger(__name__)

28 +_UNKNOWN = object()

29 +"""Sentinel meta value for proxies whose concrete value could not be inferred.

30 +Distinct from `None`, which is a valid concrete value (e.g. `attn_weights`)."""

26 -def _infer_len(node: fx.Node) -> int | None:

27 - """Concrete length of a proxy's value, inferred from its node chain.

32 +_MODULE_CALL = nn.Module.__call__

33 +"""The unpatched `nn.Module.__call__`. During tracing fx patches it to record

34 +`call_module` nodes; meta execution must call modules for real."""

35 +

36 +

37 +def is_leaf_call(node: object) -> bool:

38 + """Is node a call recorded by `_as_leaf_call` (e.g. an attention interface)."""

39 + return isinstance(node, fx.Node) and node.meta.get("leaf_call", False)

29 - Lets tracing pass through the shape unpacks and `*`-splats (e.g.

30 - `(*input_shape, -1, head_dim)`) that precede the patterns in HF attention.

31 - """

32 - # `x.shape` has the rank of `x`, when known

33 - if (

34 - node.op == "call_function"

35 - and node.target is getattr

36 - and node.args[1] == "shape"

37 - and (rank := _rank(node.args[0])) is not None

38 - ):

39 - return rank

40 - # Slices of known-length values

41 - if node.op == "call_function" and node.target is operator.getitem:

42 - src_len = _infer_len(node.args[0])

43 - index = node.args[1]

44 - if src_len is not None and isinstance(index, slice):

45 - return len(range(*index.indices(src_len)))

46 - return None

42 +def _reference_weight(module: nn.Module) -> torch.Tensor | None:

43 + """A weight whose trailing dim is the module's hidden size.

49 -def _rank(node: fx.Node) -> int | None:

50 - """The tensor rank of `node`'s value, if known."""

51 - # vLLM always feeds the model [1, seq_len, hidden_size] hidden states

52 - if node.op == "placeholder" and node.target == "hidden_states":

53 - return 3

45 + Linears and 2-D gate weights are `[out, hidden]`; norm weights are

46 + `[hidden]`. Used to fabricate a placeholder input of matching size/dtype."""

47 + for child in module.modules():

48 + if isinstance(child, nn.Linear):

49 + return child.weight

50 + for param in module.parameters():

51 + if param.ndim in (1, 2):

52 + return param

53 return None

57 -class _SizedProxy(fx.Proxy):

58 - """Proxy whose `len` is inferred from the graph (see `_infer_len`)."""

56 +class _MetaProxy(fx.Proxy):

57 + """Proxy carrying the meta-tensor value of the traced expression.

58 +

59 + Shape questions (`len`, iteration, `.shape` unpacks) are answered by

60 + executing each op on the meta values, so PyTorch's meta kernels are the

61 + single source of shape inference — no per-op rules."""

62 +

63 + meta: object = _UNKNOWN

65 def __len__(self) -> int:

61 - length = _infer_len(self.node)

62 - if length is None:

63 - return super().__len__()

64 - return length

66 + if self.meta is not _UNKNOWN:

67 + return len(self.meta)

68 + return super().__len__() # type: ignore[misc]

69 +

70 + def __getattr__(self, k: str) -> "_MetaAttribute":

71 + return _MetaAttribute(self, k)

72 +

73 +

74 +class _MetaAttribute(_MetaProxy, fx.proxy.Attribute):

75 + """Attribute proxy (e.g. `x.shape`) carrying its meta value.

76 +

77 + `Proxy.__getattr__` constructs `Attribute` directly, bypassing

78 + `Tracer.proxy`, so the meta value must be grafted on here too."""

79 +

80 + def __init__(self, root: fx.Proxy, attr: str):

81 + super().__init__(root, attr)

82 + root_meta = getattr(root, "meta", _UNKNOWN)

83 + if root_meta is not _UNKNOWN:

84 + with contextlib.suppress(Exception):

85 + self.meta = getattr(root_meta, attr)

88 class _AllLeafTracer(fx.Tracer):

@@ -69,31 +90,159 @@ class _AllLeafTracer(fx.Tracer):

91 Each child stays one `call_module` node, so matching sees the module's own

92 forward structure (activations aren't decomposed into e.g. `sigmoid * x`).

72 - `iter` traces through the leading shape unpacks (see `_infer_len`); anything

73 - else untraceable ends the trace early and the partial graph is matched.

93 + Every traced op is also executed on meta tensors (see `_MetaProxy`) so

94 + shape unpacks and `*`-splats trace through; anything else untraceable ends

95 + the trace early and the partial graph is matched.

96 """

98 + varkw: str | None = None

99 + """Name of the traced forward's `**kwargs` parameter, if any."""

100 +

101 def is_leaf_module(self, m: nn.Module, module_qualified_name: str) -> bool:

102 return True

103

104 def proxy(self, node: fx.Node) -> fx.Proxy:

80 - return _SizedProxy(node, self)

105 + return _MetaProxy(node, self)

106 +

107 + def create_proxy(self, kind, target, args, kwargs, *extra, **extra_kwargs):

108 + proxy = super().create_proxy(kind, target, args, kwargs, *extra, **extra_kwargs)

109 + if isinstance(proxy, _MetaProxy) and proxy.meta is _UNKNOWN:

110 + # A failure stays _UNKNOWN; only fatal if a shape question is asked.

111 + with contextlib.suppress(Exception):

112 + proxy.meta = self._infer_meta(kind, target, args, kwargs)

113 + return proxy

114 +

115 + def _infer_meta(self, kind: str, target: object, args: tuple, kwargs: dict):

116 + """Execute the op on meta tensors; PyTorch infers the output value."""

117 + if kind == "placeholder":

118 + # vLLM always feeds the model [1, seq_len, hidden_size] hidden states.

119 + weight = _reference_weight(self.root)

120 + if str(target) == "hidden_states" and weight is not None:

121 + return torch.empty(

122 + 1, 8, weight.shape[-1], dtype=weight.dtype, device="meta"

123 + )

124 + return _UNKNOWN

125 + if kind == "get_attr":

126 + value = operator.attrgetter(str(target))(self.root)

127 + if isinstance(value, torch.Tensor):

128 + value = torch.empty_like(value, device="meta")

129 + return value

130 + unknown = False

131 +

132 + def meta_of(arg: object) -> object:

133 + nonlocal unknown

134 + if isinstance(arg, fx.Proxy):

135 + meta = getattr(arg, "meta", _UNKNOWN)

136 + unknown = unknown or meta is _UNKNOWN

137 + return meta

138 + return arg

139 +

140 + meta_args = fx.node.map_aggregate(args, meta_of)

141 + meta_kwargs = fx.node.map_aggregate(kwargs, meta_of)

142 + if unknown:

143 + return _UNKNOWN

144 + if kind == "call_function":

145 + return target(*meta_args, **meta_kwargs)

146 + if kind == "call_method":

147 + receiver, *rest = meta_args

148 + return getattr(receiver, str(target))(*rest, **meta_kwargs)

149 + if kind == "call_module":

150 + # Run the child's forward with all its state on "meta", without

151 + # mutating it (at match time params may be meta but buffers real).

152 + # fx patches `nn.Module.__call__` while tracing; restore the real

153 + # one so this execution is not itself recorded.

154 + child = self.root.get_submodule(str(target))

155 + state = {

156 + name: torch.empty_like(tensor, device="meta")

157 + for name, tensor in chain(

158 + child.named_parameters(), child.named_buffers()

159 + )

160 + }

161 + with mock.patch.object(nn.Module, "__call__", _MODULE_CALL):

162 + return torch.func.functional_call(child, state, meta_args, meta_kwargs)

163 + return _UNKNOWN

164 +

165 + def _is_varkw(self, node: object) -> bool:

166 + return (

167 + isinstance(node, fx.Node)

168 + and node.op == "placeholder"

169 + and str(node.target).lstrip("*") == self.varkw

170 + )

171

172 def iter(self, obj: fx.Proxy):

83 - length = _infer_len(obj.node)

84 - if length is None:

173 + # Assume kwargs is always empty to simplify tracing.

174 + node = obj.node

175 + if self._is_varkw(node) or (

176 + node.op == "call_method"

177 + and node.target == "keys"

178 + and self._is_varkw(node.args[0])

179 + ):

180 + return iter(())

181 + meta = getattr(obj, "meta", _UNKNOWN)

182 + if meta is _UNKNOWN:

183 return super().iter(obj)

86 - return iter([obj[i] for i in range(length)])

184 + return iter([obj[i] for i in range(len(meta))])

185

186

89 -def trace(module: nn.Module) -> fx.Graph | None:

90 - """Trace `module.forward`, returning the partial graph on failure.

187 +def _as_leaf_call(fn: Callable, length: int | None = None) -> Callable:

188 + """Wrap any callable so tracing records it as one opaque `call_function` node.

189 +

190 + Lets the trace continue past untraceable bodies. Only the proxy arguments carry into

191 + the node's dataflow; the rest are dropped rather than lifted into the graph.

192 + `length` declares how many values the callable returns, so unpacking its result also

193 + traces. Called without proxies (i.e. outside tracing), the wrapper is a passthrough.

194 + """

195 +

196 + def leaf(*args, **kwargs):

197 + proxies = tuple(arg for arg in args if isinstance(arg, fx.Proxy))

198 + if not proxies:

199 + return fn(*args, **kwargs)

200 + proxy = proxies[0].tracer.create_proxy("call_function", fn, proxies, {})

201 + proxy.node.meta["leaf_call"] = True

202 + if length is not None:

203 + # The body never executes, so fabricate a value of the declared length.

204 + proxy.meta = (_UNKNOWN,) * length

205 + return proxy

206 +

207 + return leaf

208 +

209 +

210 +def _leaf_attention_interfaces():

211 + """Patch `AttentionInterface.get_interface` so traced forwards see a leaf node.

212

92 - The graph is only evidence for matching, and the patterns sit at the top of

93 - their forwards, so a trace that fails partway can still be matched."""

213 + `vllm_attention_function` needs runtime context so it is untraceable.

214 + Every interface returns `(attn_output, attn_weights)`."""

215 + from transformers.modeling_utils import AttentionInterface

216 +

217 + original = AttentionInterface.get_interface

218 +

219 + def get_interface(self, *args, **kwargs):

220 + return _as_leaf_call(original(self, *args, **kwargs), length=2)

221 +

222 + return mock.patch.object(AttentionInterface, "get_interface", get_interface)

223 +

224 +

225 +def trace(module: nn.Module) -> fx.Graph | None:

226 + """Trace `module.forward`, returning the partial graph on failure."""

227 + parameters = forward_parameters(type(module))

228 + # vLLM never passes `past_key_values` so it is always the default value of `None`.

229 + # Make this concrete to simplify tracing.

230 + concrete_args = None

231 + if "past_key_values" in parameters:

232 + concrete_args = {"past_key_values": None}

233 + # Get the name of the kwargs parameter passed to module.forward (usually "kwargs")

234 tracer = _AllLeafTracer()

235 + tracer.varkw = next(

236 + (

237 + p.name

238 + for p in parameters.values()

239 + if p.kind is inspect.Parameter.VAR_KEYWORD

240 + ),

241 + None,

242 + )

243 try:

96 - return tracer.trace(module)

244 + with _leaf_attention_interfaces():

245 + return tracer.trace(module, concrete_args=concrete_args)

246 except Exception as exc:

247 logger.debug("Could not fully trace %s: %s", type(module), exc)

248 return getattr(tracer, "graph", None)

@@ -129,20 +278,27 @@ def recover_forward(cls: type[nn.Module]) -> tuple[ast.FunctionDef, Callable]:

278 return funcdef, fn

279

280

281 +def forward_parameters(cls: type[nn.Module]) -> dict[str, inspect.Parameter]:

282 + """`cls.forward`'s signature parameters, or empty if uninspectable."""

283 + try:

284 + return dict(inspect.signature(cls.forward).parameters)

285 + except (TypeError, ValueError):

286 + return {}

287 +

288 +

289 def forward_input_count(cls: type[nn.Module]) -> int:

290 """The number of tensor inputs `cls.forward` declares, excluding `self` and

291 any `*args`/`**kwargs`. Read from the signature, so it is independent of

292 whether the trace completes (unlike counting placeholders)."""

136 - try:

137 - params = list(inspect.signature(cls.forward).parameters.values())[1:]

138 - except (ValueError, TypeError):

293 + params = list(forward_parameters(cls).values())

294 + if not params:

295 return 1 # uninspectable: assume a single input and let matching decide

296 fixed = (

297 inspect.Parameter.POSITIONAL_ONLY,

298 inspect.Parameter.POSITIONAL_OR_KEYWORD,

299 inspect.Parameter.KEYWORD_ONLY,

300 )

145 - return sum(1 for p in params if p.kind in fixed)

301 + return sum(1 for p in params[1:] if p.kind in fixed)

302

303

304 def compile_forward(funcdef: ast.FunctionDef, fn: Callable) -> Callable:

@@ -227,6 +383,42 @@ def find_node(graph: fx.Graph, predicate: Callable[[fx.Node], bool]) -> fx.Node

383 return next((n for n in graph.nodes if predicate(n)), None)

384

385

386 +def output_value(graph: fx.Graph) -> object | None:

387 + """The value the graph's `output` node returns, if the trace reached one."""

388 + output = find_node(graph, lambda n: n.op == "output")

389 + if output is None or not output.args:

390 + return None

391 + return output.args[0]

392 +

393 +

394 +def upstream_linear(node: object, module: nn.Module) -> fx.Node | None:

395 + """Nearest linear producing `node`, walking back through splits/reshapes.

396 +

397 + Never walks through a leaf call (e.g. an attention interface): its inputs

398 + are what attention consumes, not what produced the value."""

399 + stack = [node]

400 + seen: set[fx.Node] = set()

401 + while stack:

402 + current = stack.pop()

403 + if not isinstance(current, fx.Node) or current in seen:

404 + continue

405 + seen.add(current)

406 + if is_linear(current, module):

407 + return current

408 + if current.op in ("call_function", "call_method") and not is_leaf_call(current):

409 + stack.extend(current.args)

410 + return None

411 +

412 +

413 +def returned_linear(graph: fx.Graph, module: nn.Module) -> str | None:

414 + """Name of the Linear producing the graph's (first) output value."""

415 + value = output_value(graph)

416 + if isinstance(value, (tuple, list)) and value:

417 + value = value[0]

418 + linear = upstream_linear(value, module)

419 + return None if linear is None else str(linear.target)

420 +

421 +

422 def is_linear(node: fx.Node, module: nn.Module) -> bool:

423 """Is node `nn.Linear.__call__()`."""

424 return node.op == "call_module" and isinstance(