wehub-resource-sync
259 lines
8.7 KiB
Python
259 lines
8.7 KiB
Python
#
|
|
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
#
|
|
|
|
"""
|
|
Construct a self-attention ONNX model using the ONNX GraphSurgeon layer API.
|
|
|
|
The model implements:
|
|
1. Q/K/V linear projections (MatMul with 4096x4096 weights)
|
|
2. Reshape to multi-head layout (seq, batch, 4096) -> (seq, batch, 32, 128)
|
|
3. RMSNorm on Q and K
|
|
4. Scaled dot-product attention (QK^T / sqrt(d), softmax, attn * V)
|
|
5. Reshape back and output projection
|
|
|
|
Input/Output: (sequence_length, batch_size, 4096), float16
|
|
"""
|
|
|
|
import argparse
|
|
import math
|
|
|
|
import numpy as np
|
|
import onnx
|
|
import onnx_graphsurgeon as gs
|
|
|
|
NUM_HEADS = 32
|
|
HEAD_DIM = 128
|
|
HIDDEN_DIM = NUM_HEADS * HEAD_DIM # 4096
|
|
OPSET = 17
|
|
|
|
|
|
# Register ONNX ops as methods on gs.Graph using the layer API.
|
|
# Each returns the output tensor(s) directly for easy chaining.
|
|
|
|
@gs.Graph.register()
|
|
def matmul(self, a, b):
|
|
return self.layer(op="MatMul", inputs=[a, b], outputs=["matmul_out"])[0]
|
|
|
|
|
|
@gs.Graph.register()
|
|
def transpose(self, a, perm):
|
|
return self.layer(op="Transpose", inputs=[a], attrs={"perm": perm}, outputs=["transpose_out"])[0]
|
|
|
|
|
|
@gs.Graph.register()
|
|
def reshape(self, data, shape):
|
|
return self.layer(op="Reshape", inputs=[data, shape], attrs={"allowzero": 0}, outputs=["reshape_out"])[0]
|
|
|
|
|
|
@gs.Graph.register()
|
|
def softmax(self, a, axis=-1):
|
|
return self.layer(op="Softmax", inputs=[a], attrs={"axis": axis}, outputs=["softmax_out"])[0]
|
|
|
|
|
|
@gs.Graph.register()
|
|
def cast(self, a, to):
|
|
return self.layer(op="Cast", inputs=[a], attrs={"to": to}, outputs=["cast_out"])[0]
|
|
|
|
|
|
@gs.Graph.register()
|
|
def sqrt(self, a):
|
|
return self.layer(op="Sqrt", inputs=[a], outputs=["sqrt_out"])[0]
|
|
|
|
|
|
@gs.Graph.register()
|
|
def add(self, a, b):
|
|
return self.layer(op="Add", inputs=[a, b], outputs=["add_out"])[0]
|
|
|
|
|
|
@gs.Graph.register()
|
|
def mul(self, a, b):
|
|
return self.layer(op="Mul", inputs=[a, b], outputs=["mul_out"])[0]
|
|
|
|
|
|
@gs.Graph.register()
|
|
def div(self, a, b):
|
|
return self.layer(op="Div", inputs=[a, b], outputs=["div_out"])[0]
|
|
|
|
|
|
@gs.Graph.register()
|
|
def pow(self, a, b):
|
|
return self.layer(op="Pow", inputs=[a, b], outputs=["pow_out"])[0]
|
|
|
|
|
|
@gs.Graph.register()
|
|
def reduce_mean(self, a, axes, keepdims=1):
|
|
return self.layer(
|
|
op="ReduceMean", inputs=[a], attrs={"axes": axes, "keepdims": keepdims},
|
|
outputs=["reduce_mean_out"],
|
|
)[0]
|
|
|
|
|
|
@gs.Graph.register()
|
|
def shape_op(self, a):
|
|
return self.layer(op="Shape", inputs=[a], outputs=["shape_out"])[0]
|
|
|
|
|
|
@gs.Graph.register()
|
|
def gather(self, data, indices):
|
|
return self.layer(op="Gather", inputs=[data, indices], attrs={"axis": 0}, outputs=["gather_out"])[0]
|
|
|
|
|
|
@gs.Graph.register()
|
|
def unsqueeze(self, a, axes):
|
|
return self.layer(op="Unsqueeze", inputs=[a, axes], outputs=["unsqueeze_out"])[0]
|
|
|
|
|
|
@gs.Graph.register()
|
|
def concat(self, inputs, axis=0):
|
|
return self.layer(op="Concat", inputs=inputs, attrs={"axis": axis}, outputs=["concat_out"])[0]
|
|
|
|
|
|
def build_attention_graph():
|
|
"""Build the full self-attention ONNX graph."""
|
|
rng = np.random.default_rng(42)
|
|
graph = gs.Graph(opset=OPSET)
|
|
|
|
def fp16_weights(shape):
|
|
return rng.standard_normal(shape).astype(np.float16)
|
|
|
|
def fp32_scalar(val):
|
|
return np.array([val], dtype=np.float32)
|
|
|
|
axes_0 = np.array([0], dtype=np.int64)
|
|
|
|
# Input: (seq, batch, 4096) fp16
|
|
graph_input = gs.Variable("input", dtype=np.float16, shape=["sequence_length", "batch_size", HIDDEN_DIM])
|
|
graph.inputs = [graph_input]
|
|
|
|
# Q/K/V projections
|
|
q_proj = graph.matmul(graph_input, fp16_weights((HIDDEN_DIM, HIDDEN_DIM)))
|
|
k_proj = graph.matmul(graph_input, fp16_weights((HIDDEN_DIM, HIDDEN_DIM)))
|
|
v_proj = graph.matmul(graph_input, fp16_weights((HIDDEN_DIM, HIDDEN_DIM)))
|
|
|
|
# Dynamic reshape: (seq, batch, 4096) -> (seq, batch, 32, 128)
|
|
# Build target shape [seq_dim, batch_dim, 32, 128] from input shape
|
|
def reshape_to_heads(proj):
|
|
inp_shape = graph.shape_op(proj)
|
|
seq_dim = graph.unsqueeze(graph.gather(inp_shape, np.array(0, dtype=np.int64)), axes_0)
|
|
batch_dim = graph.unsqueeze(graph.gather(inp_shape, np.array(1, dtype=np.int64)), axes_0)
|
|
target_shape = graph.concat([
|
|
seq_dim, batch_dim,
|
|
np.array([NUM_HEADS], dtype=np.int64),
|
|
np.array([HEAD_DIM], dtype=np.int64),
|
|
])
|
|
return graph.reshape(proj, target_shape)
|
|
|
|
q_4d = reshape_to_heads(q_proj)
|
|
k_4d = reshape_to_heads(k_proj)
|
|
v_4d = reshape_to_heads(v_proj)
|
|
|
|
# RMSNorm: x * rsqrt(mean(x^2) + eps) * weight
|
|
def rmsnorm(x):
|
|
x_fp32 = graph.cast(x, onnx.TensorProto.FLOAT)
|
|
sq = graph.pow(x_fp32, fp32_scalar(2.0))
|
|
mean = graph.reduce_mean(sq, axes=[-1])
|
|
rms = graph.sqrt(graph.add(mean, fp32_scalar(1e-6)))
|
|
inv_rms = graph.div(fp32_scalar(1.0), rms)
|
|
normed = graph.mul(x_fp32, inv_rms)
|
|
normed_fp16 = graph.cast(normed, onnx.TensorProto.FLOAT16)
|
|
weight = rng.standard_normal((1, 1, 1, HEAD_DIM)).astype(np.float16)
|
|
return graph.mul(weight, normed_fp16)
|
|
|
|
q_norm = rmsnorm(q_4d)
|
|
k_norm = rmsnorm(k_4d)
|
|
|
|
# Transpose to attention layout: (seq, batch, heads, hdim) -> (batch, heads, seq, hdim)
|
|
q_attn = graph.transpose(q_norm, perm=[1, 2, 0, 3])
|
|
k_attn = graph.transpose(k_norm, perm=[1, 2, 0, 3])
|
|
v_attn = graph.transpose(v_4d, perm=[1, 2, 0, 3])
|
|
|
|
# Dynamic reshape Q/K/V to (batch, heads, -1, hdim) using shape extraction
|
|
def reshape_attn(x):
|
|
s = graph.shape_op(x)
|
|
batch = graph.unsqueeze(graph.gather(s, np.array(0, dtype=np.int64)), axes_0)
|
|
heads = graph.unsqueeze(graph.gather(s, np.array(1, dtype=np.int64)), axes_0)
|
|
hdim = graph.unsqueeze(graph.gather(s, np.array(3, dtype=np.int64)), axes_0)
|
|
target = graph.concat([batch, heads, np.array([-1], dtype=np.int64), hdim])
|
|
return graph.reshape(x, target)
|
|
|
|
q_r = reshape_attn(q_attn)
|
|
k_r = reshape_attn(k_attn)
|
|
v_r = reshape_attn(v_attn)
|
|
|
|
# Scale: split sqrt(1/sqrt(head_dim)) across Q and K
|
|
scale_val = math.sqrt(math.sqrt(1.0 / HEAD_DIM))
|
|
scale_fp16 = np.array([scale_val], dtype=np.float16)
|
|
q_scaled = graph.mul(q_r, scale_fp16)
|
|
q_scaled.name = "q_scaled"
|
|
|
|
# Transpose K: (batch, heads, seq, hdim) -> (batch, heads, hdim, seq)
|
|
k_t = graph.transpose(k_r, perm=[0, 1, 3, 2])
|
|
k_scaled = graph.mul(k_t, scale_fp16)
|
|
|
|
# QK^T -> Softmax -> Attn*V
|
|
qk = graph.matmul(q_scaled, k_scaled)
|
|
attn_weights = graph.softmax(qk, axis=-1)
|
|
attn_out = graph.matmul(attn_weights, v_r)
|
|
|
|
# Reshape back: (batch, heads, seq, hdim) -> (seq, batch, 4096)
|
|
attn_t = graph.transpose(attn_out, perm=[2, 0, 1, 3])
|
|
attn_shape = graph.shape_op(attn_t)
|
|
seq_dim = graph.unsqueeze(graph.gather(attn_shape, np.array(0, dtype=np.int64)), axes_0)
|
|
batch_dim = graph.unsqueeze(graph.gather(attn_shape, np.array(1, dtype=np.int64)), axes_0)
|
|
# Compute heads * hdim
|
|
heads_dim = graph.gather(attn_shape, np.array(2, dtype=np.int64))
|
|
hdim_dim = graph.gather(attn_shape, np.array(3, dtype=np.int64))
|
|
hidden = graph.unsqueeze(graph.mul(heads_dim, hdim_dim), axes_0)
|
|
flat_shape = graph.concat([seq_dim, batch_dim, hidden])
|
|
attn_flat = graph.reshape(attn_t, flat_shape)
|
|
|
|
# Output projection
|
|
output = graph.matmul(attn_flat, fp16_weights((HIDDEN_DIM, HIDDEN_DIM)))
|
|
output.name = "output"
|
|
output.dtype = np.float16
|
|
output.shape = ["sequence_length", "batch_size", HIDDEN_DIM]
|
|
graph.outputs = [output]
|
|
|
|
graph.cleanup().toposort()
|
|
model = gs.export_onnx(graph)
|
|
model.ir_version = 8
|
|
|
|
return model
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
description="Generate attention_sd.onnx model using the ONNX GraphSurgeon layer API"
|
|
)
|
|
parser.add_argument(
|
|
"--output", type=str, default="attention_sd.onnx",
|
|
help="Output ONNX file path (default: attention_sd.onnx)",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
model = build_attention_graph()
|
|
onnx.save(model, args.output)
|
|
|
|
print(f"Saved model to {args.output}")
|
|
print(f" Nodes: {len(model.graph.node)}")
|
|
print(f" Initializers: {len(model.graph.initializer)}")
|
|
print(f" Opset: {model.opset_import[0].version}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|