@dataclass
class QKVFuser(StackedFuser):
"""Fuser for the attention QKV pattern `q(x), k(x), v(x)`."""
q_name: str
k_name: str
v_name: str
o_name: str | None
merged_name: ClassVar[str] = "qkv_proj"
merged_cls: ClassVar[str] = "QKVParallelLinear"
@property
def shards(self) -> list[tuple[str, ShardId]]:
return [(self.q_name, "q"), (self.k_name, "k"), (self.v_name, "v")]
@classmethod
def _get_qkv_nodes(
cls, graph: fx.Graph, module: nn.Module
) -> tuple[fx.Node, fx.Node, fx.Node] | None:
"""Search `graph` for the QKV pattern `q(x), k(x), v(x)`."""
by_input: dict[fx.Node, list[fx.Node]] = {}
for node in graph.nodes:
if (
is_linear(node, module)
and len(node.args) == 1
and not node.kwargs
and isinstance(node.args[0], fx.Node)
and node.args[0].op == "placeholder"
):
by_input.setdefault(node.args[0], []).append(node)
triples = [nodes for nodes in by_input.values() if len(nodes) == 3]
if len(triples) != 1:
return None
q_node, k_node, v_node = nodes = triples[0]
outs = [module.get_submodule(node.target).out_features for node in nodes]
if len(set(outs)) == 2:
# q is identified as the larger projection (GQA)
(q_node,) = (n for n, out in zip(nodes, outs) if outs.count(out) == 1)
k_node, v_node = (n for n, out in zip(nodes, outs) if outs.count(out) == 2)
if module.get_submodule(q_node.target).out_features != max(outs):
return None
elif len(set(outs)) != 1:
return None
return q_node, k_node, v_node
@classmethod
def match(cls, graph: fx.Graph, module: nn.Module) -> "QKVFuser | None":
if (qkv_nodes := cls._get_qkv_nodes(graph, module)) is None:
return None
q, k, v = qkv_nodes
names = dict(q_name=q.target, k_name=k.target, v_name=v.target)
attn_width = module.get_submodule(q.target).out_features
candidates = [
name
for name, child in module.named_children()
if isinstance(child, nn.Linear)
and name not in names.values()
and child.in_features == attn_width
]
names["o_name"] = candidates[0] if len(candidates) == 1 else None
return cls(source_cls=type(module).__name__, **names)
def update_forward(self, module: nn.Module) -> None:
"""Replace `q(x), k(x), v(x)` with `qkv(x).split(sizes, -1)` in source."""
funcdef, fn = recover_forward(type(module))
calls = [
single_self_call(funcdef, name)
for name in (self.q_name, self.k_name, self.v_name)
]
arg_dumps = {ast.dump(call.args[0]) for call in calls}
if len(arg_dumps) != 1:
raise ValueError("projection inputs are written differently")
# The trace may be partial, so prove projection exclusivity in source:
# no other linear child may consume the same input (else the matched
# three may not be q, k and v)
other_linears = {
name
for name, child in module.named_children()
if isinstance(child, nn.Linear)
} - {self.q_name, self.k_name, self.v_name}
for node in ast.walk(funcdef):
if (
isinstance(node, ast.Call)
and isinstance(node.func, ast.Attribute)
and node.func.attr in other_linears
and any(ast.dump(arg) in arg_dumps for arg in node.args)
):
raise ValueError("another linear consumes the same input")
blocks = [innermost_block(funcdef.body, call) for call in calls]
if any(found is None for found in blocks):
raise ValueError("projection calls not found in the function body")
if len({id(block) for block, _ in blocks}) != 1:
raise ValueError("projection calls are in different blocks")
# q(x), k(x), v(x) -> q, k, v = qkv(x).split(qkv.output_sizes / qkv.tp_size, -1)
names = {node.id for node in ast.walk(funcdef) if isinstance(node, ast.Name)}
temps = [f"{name}_fused" for name in (self.q_name, self.k_name, self.v_name)]
if names & set(temps):
raise ValueError("fused temporaries would shadow existing names")
merged = f"self.{self.merged_name}"
sections = f"[s // {merged}.tp_size for s in {merged}.output_sizes]"
template = f"{', '.join(temps)} = {merged}(__arg__).split({sections}, -1)"
assign = ast.parse(template).body[0]
arg = next(
node
for node in ast.walk(assign)
if isinstance(node, ast.Name) and node.id == "__arg__"
)
replace_expr(assign, arg, calls[0].args[0])
block, index = blocks[0]
ast.copy_location(assign, block[index])
block.insert(min(index for _, index in blocks), assign)
for call, temp in zip(calls, temps):
replace_expr(funcdef, call, ast.Name(id=temp, ctx=ast.Load()))
self.fused_forward = compile_forward(funcdef, fn)
def validate(self, module: nn.Module, model_config: "ModelConfig") -> bool:
"""Shapes must be compatible for a single merged, head-sharded GEMM."""
q = module.get_submodule(self.q_name)
k = module.get_submodule(self.k_name)
v = module.get_submodule(self.v_name)
head_size = model_config.get_head_size()
compatible = (
q.in_features == k.in_features == v.in_features
and len({proj.bias is None for proj in (q, k, v)}) == 1
and k.out_features == v.out_features
and q.out_features % head_size == 0
and k.out_features % head_size == 0
)
if not compatible:
logger.debug("%s is not compatible with QKV fusion", type(module))
return compatible
def update_attrs(
self,
module: nn.Module,
prefix: str,
model_config: "ModelConfig",
quant_config: "QuantizationConfig",
) -> None:
head_size = model_config.get_head_size()
q = module.get_submodule(self.q_name)
k = module.get_submodule(self.k_name)
merged = QKVParallelLinear(
hidden_size=q.in_features,
head_size=head_size,
total_num_heads=q.out_features // head_size,
total_num_kv_heads=k.out_features // head_size,
bias=q.bias is not None,
quant_config=quant_config,
prefix=maybe_prefix(prefix, self.merged_name),
return_bias=False,
)
logger.debug(
"%s: %s, %s: %s, %s: %s -> %s: %s",
self.q_name,
q,
self.k_name,
k,
self.v_name,
module.get_submodule(self.v_name),
self.merged_name,
merged,
)
setattr(module, self.merged_name, merged)
# Drop the consumed submodules so their (meta) params are not expected.
for name in (self.q_name, self.k_name, self.v_name):
delattr(module, name)
# If there is an output projection, we know it must be rowwise.
if self.o_name is not None:
o_proj_prefix = maybe_prefix(prefix, self.o_name)
o_proj = module.get_submodule(self.o_name)
new_o = replace_linear_class(
o_proj, "rowwise", quant_config, prefix=o_proj_prefix
)
setattr(module, self.o_name, new_o)
log_replacement(o_proj_prefix, o_proj, new_o)