wehub/apache--tvm
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
Lint / lint (push) Has been cancelled
Details
CI / MacOS (push) Has been cancelled
Details
CI / Windows (push) Has been cancelled
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:36:25 +08:00
commit 26446540fa
3151 changed files with 974126 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/python/relax/backend/adreno/mod_utils.py
+842
View File
@@ -0,0 +1,842 @@
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you 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.
# ruff: noqa: E501, F401, F841
"""CLML integration operator tests."""
import json
import numpy as np
import pytest
import tvm
import tvm.testing
from tvm import relax, rpc
from tvm.relax.backend.adreno import clml
from tvm.script import ir as I
from tvm.script import relax as R
from tvm.script import tirx as T
from tvm.script.ir_builder import IRBuilder
from tvm.script.ir_builder import relax as relax_builder
def get_relax_conv2d_mod(
data_shape,
weight_shape,
stride,
dilation,
padding,
weight_layout="OIHW",
groups=1,
dtype="float32",
has_bias=False,
has_bn=False,
has_activation=False,
has_pad=False,
is_depthwise=False,
):
with IRBuilder() as builder:
with relax_builder.function():
R.func_name("main")
if has_pad:
p = (0, 0, 0, 0, padding[0], padding[0], padding[1], padding[1])
orig_data = R.arg("data", R.Tensor(data_shape, dtype))
data = R.nn.pad(orig_data, pad_width=p, pad_value=0.0)
padding = (0, 0, 0, 0)
else:
data = R.arg("data", R.Tensor(data_shape, dtype))
weight = R.arg("weight", R.Tensor(weight_shape, dtype))
if has_bias:
bias = R.arg("bias", R.Tensor((1, weight_shape[0], 1, 1), dtype))
is_depthwise = data_shape[1] == weight_shape[0] == groups
with R.dataflow() as frame:
output = R.emit(
R.nn.conv2d(
data,
weight,
out_dtype=dtype,
strides=stride,
dilation=dilation,
padding=padding,
data_layout="NCHW",
kernel_layout=weight_layout,
groups=groups,
)
)
if has_bias:
output = R.emit(output + bias)
if has_bn:
gamma = R.arg("gamma", R.Tensor((weight_shape[0],), dtype))
beta = R.arg("beta", R.Tensor((weight_shape[0],), dtype))
mean = R.arg("mean", R.Tensor((weight_shape[0],), dtype))
variance = R.arg("variance", R.Tensor((weight_shape[0],), dtype))
output = R.emit(
R.nn.batch_norm(output, gamma, beta, mean, variance, axis=1, epsilon=1e-5)[
0
]
)
if has_activation:
output = R.emit(R.nn.relu(output))
R.output(output)
R.func_ret_value(frame.output_vars[0])
func = builder.get()
return tvm.IRModule({"main": func})
def get_clml_conv2d_codegen(
data_shape,
weight_shape,
stride,
dilation,
padding,
weight_layout="OIHW",
groups=1,
dtype="float32",
has_bias=False,
has_bn=False,
has_activation=False,
has_pad=False,
is_depthwise=False,
):
kernel_h, kernel_w = weight_shape[2], weight_shape[3]
channels = weight_shape[0]
if len(padding) == 2:
padding = (padding[0], padding[1], padding[0], padding[1])
output_height = ((data_shape[2] - kernel_h + padding[0] + padding[2]) / stride[0]) + 1
output_width = ((data_shape[3] - kernel_w + padding[1] + padding[3]) / stride[1]) + 1
output_shape = (1, channels, int(output_height), int(output_width))
out_dtype = dtype
is_depthwise = data_shape[1] == channels == groups
weight_layout = "IOHW" if is_depthwise else "OIHW"
if weight_layout == "OIHW":
weight_shape = (channels, data_shape[1] // groups, kernel_h, kernel_w)
else:
weight_shape = (data_shape[1] // groups, channels, kernel_h, kernel_w)
if is_depthwise:
name = "openclml.nn.depthwise_conv2d"
else:
name = "openclml.nn.conv2d"
node = {
"op": "kernel",
"name": "",
"inputs": [],
"attrs": {
"groups": groups,
"num_outputs": 1,
"data_layout": "NCHW",
"kernel_layout": weight_layout,
"dilation": dilation,
"out_layout": "NCHW",
"out_dtype": out_dtype,
"shape": [list(output_shape)],
"dtype": [dtype],
"padding": padding,
"strides": stride,
},
}
if has_activation:
node["attrs"]["activation_type"] = "relu"
nodes = [
{
"op": "input",
"name": "",
"attrs": {"shape": [list(data_shape)], "dtype": [str(dtype)]},
},
]
nodes.append(
{
"op": "const",
"name": "",
"attrs": {"shape": [list(weight_shape)], "dtype": [str(dtype)]},
}
)
if has_bias:
bias_dtype = dtype
nodes.append(
{
"op": "const",
"name": "",
"attrs": {
"shape": [[1, weight_shape[1] if is_depthwise else weight_shape[0], 1, 1]],
"dtype": [bias_dtype],
},
}
)
if has_bn:
bn_shape = [1, weight_shape[0], 1, 1]
# conv2d + bn --> conv2d + Add due to OptimizeBatchNorm transformation Pass
nodes.append(
{
"name": "",
"op": "const",
"attrs": {"dtype": [dtype], "shape": [[1, weight_shape[0], 1, 1]]},
},
)
input_idx = 0
for _ in range(len(nodes)):
node["inputs"].append([input_idx, 0, 0])
input_idx += 1
node["attrs"]["num_inputs"] = len(nodes)
nodes.append(node)
return nodes
def get_relax_conv2d_transpose_mod(
data_shape,
weight_shape,
channels,
stride,
padding,
dtype="float32",
):
with IRBuilder() as builder:
with relax_builder.function():
R.func_name("main")
data = R.arg("data", R.Tensor(data_shape, dtype))
weight = R.arg("weight", R.Tensor(weight_shape, dtype))
with R.dataflow() as frame:
output = R.emit(
R.nn.conv2d_transpose(
data,
weight,
groups=1,
strides=stride,
padding=padding,
kernel_layout="OIHW",
data_layout="NCHW",
)
)
R.output(output)
R.func_ret_value(frame.output_vars[0])
func = builder.get()
return tvm.IRModule({"main": func})
def get_conv2d_transpose_expected_codegen(
dshape, kshape, channels, kernel_size, strides, padding, dilation, dtype, output_shape
):
attrs = {
"data_layout": "NCHW",
"kernel_layout": "OIHW",
"groups": 1,
"dilation": dilation,
"num_inputs": 2,
"num_outputs": 1,
"padding": padding,
"shape": [list(output_shape)],
"dtype": [dtype],
"strides": strides,
"out_dtype": "",
"out_layout": "NCHW",
"output_padding": [0, 0],
}
exp_codegen = [
{
"op": "input",
"name": "",
"attrs": {"shape": [list(dshape)], "dtype": [str(dtype)]},
},
{
"op": "const",
"name": "",
"attrs": {"shape": [list(kshape)], "dtype": [str(dtype)]},
},
{
"op": "kernel",
"name": "",
"inputs": [[0, 0, 0], [1, 0, 0]],
"attrs": attrs,
},
]
return exp_codegen
def get_batchnorm_mod(data_shape, channels, axis, epsilon, dtype):
with IRBuilder() as builder:
with relax_builder.function():
R.func_name("main")
data = R.arg("data", R.Tensor(data_shape, dtype))
gamma = R.arg("gamma", R.Tensor((channels,), dtype))
beta = R.arg("beta", R.Tensor((channels,), dtype))
mean = R.arg("moving_mean", R.Tensor((channels,), dtype))
variance = R.arg("moving_var", R.Tensor((channels,), dtype))
with R.dataflow() as frame:
output = R.emit(
R.nn.batch_norm(data, gamma, beta, mean, variance, axis, epsilon)[0]
)
R.output(output)
R.func_ret_value(frame.output_vars[0])
func = builder.get()
return tvm.IRModule({"main": func})
def get_binary_op_mod(a_shape, b_shape, op, dtype):
with IRBuilder() as builder:
with relax_builder.function():
R.func_name("main")
a = R.arg("a", R.Tensor(a_shape, dtype))
b = R.arg("b", R.Tensor(b_shape, dtype))
with R.dataflow() as frame:
output = R.emit(op(a, b))
R.output(output)
R.func_ret_value(frame.output_vars[0])
func = builder.get()
low, high = 0, 1
a_data = np.random.uniform(low, high, size=(a_shape)).astype(dtype)
b_data = np.random.uniform(low, high, size=(b_shape)).astype(dtype)
return (tvm.IRModule({"main": func}), (a_data, b_data))
def get_unary_op_mod(a_shape, op, dtype):
with IRBuilder() as builder:
with relax_builder.function():
R.func_name("main")
a = R.arg("a", R.Tensor(a_shape, dtype))
with R.dataflow() as frame:
output = R.emit(op(a))
R.output(output)
R.func_ret_value(frame.output_vars[0])
func = builder.get()
low, high = 0, 1
a_data = np.random.uniform(low, high, size=(a_shape)).astype(dtype)
return (tvm.IRModule({"main": func}), (a_data,))
def get_relax_maxpool_mod(
data_shape, dtype, pool_size, stride=None, dilation=(1, 1), padding=(0, 0), has_pad=False
):
"""
Args:
data_shape (tuple): Input tensor shape
pool_size (tuple): Pooling window size (height, width)
stride (tuple, optional): Stride of pooling operation. Defaults to pool_size.
dilation (tuple, optional): Dilation rate. Defaults to (1, 1).
padding (tuple, optional): Padding for the input tensor. Defaults to (0, 0).
dtype (str, optional): Data type. Defaults to "float32".
has_pad (bool, optional): Whether to apply explicit padding. Defaults to False.
Returns:
tvm.IRModule: Relax MaxPool module
"""
with IRBuilder() as builder:
with relax_builder.function():
R.func_name("main")
if has_pad:
p = (0, 0, 0, 0, padding[0], padding[1], padding[0], padding[1])
orig_data = R.arg("data", R.Tensor(data_shape, dtype))
data = R.nn.pad(orig_data, pad_width=p, pad_value=float("-inf"))
padding = (0, 0)
else:
data = R.arg("data", R.Tensor(data_shape, dtype))
with R.dataflow() as frame:
output = R.emit(
R.nn.max_pool2d(
data,
pool_size=pool_size,
strides=stride,
dilation=dilation,
padding=padding,
layout="NCHW",
)
)
R.output(output)
R.func_ret_value(frame.output_vars[0])
func = builder.get()
return tvm.IRModule({"main": func})
def get_maxpool_expected_codegen(input_shape, pool_size, stride, padding, pool_type, dtype):
import math
adjusted_input_shape = [
input_shape[0],
input_shape[1],
input_shape[2] + padding[0] + padding[1],
input_shape[3] + padding[2] + padding[3],
]
pool_height = math.floor(((adjusted_input_shape[2] - pool_size[0]) / stride[0]) + 1)
pool_width = math.floor(((adjusted_input_shape[3] - pool_size[1]) / stride[1]) + 1)
output_shape = [adjusted_input_shape[0], adjusted_input_shape[1], pool_height, pool_width]
attrs = {
"ceil_mode": 0,
"dilation": [1, 1],
"layout": "NCHW",
"num_inputs": 1,
"num_outputs": 1,
"out_layout": "NCHW",
"padding": list(padding),
"pool_size": pool_size,
"shape": [list(output_shape)],
"dtype": [dtype],
"strides": stride,
"count_include_pad": 0,
}
if sum(padding):
attrs["count_include_pad"] = 0
exp_codegen = [
{
"op": "input",
"name": "",
"attrs": {"shape": [list(adjusted_input_shape)], "dtype": [str(dtype)]},
},
{
"op": "kernel",
"name": "",
"inputs": [[0, 0, 0]],
"attrs": attrs,
},
]
return exp_codegen
def get_relax_avgpool_mod(data_shape, dtype, pool_size, stride, dilation, padding, has_pad):
"""
Args:
data_shape (tuple): Input tensor shape
pool_size (tuple): Pooling window size (height, width)
stride (tuple, optional): Stride of pooling operation. Defaults to pool_size.
dilation (tuple, optional): Dilation rate. Defaults to (1, 1).
padding (tuple, optional): Padding for the input tensor. Defaults to (0, 0).
dtype (str, optional): Data type. Defaults to "float32".
has_pad (bool, optional): Whether to apply explicit padding. Defaults to False.
count_include_pad (bool, optional): Whether to include padding in averaging. Defaults to True.
Returns:
tvm.IRModule: Relax AvgPool module
"""
with IRBuilder() as builder:
with relax_builder.function():
R.func_name("main")
if has_pad:
p = (0, 0, 0, 0, padding[0], padding[1], padding[0], padding[1])
orig_data = R.arg("data", R.Tensor(data_shape, dtype))
data = R.nn.pad(orig_data, pad_width=p, pad_value=0.0)
padding = (0, 0)
else:
data = R.arg("data", R.Tensor(data_shape, dtype))
with R.dataflow() as frame:
output = R.emit(
R.nn.avg_pool2d(
data,
pool_size=pool_size,
strides=stride,
dilation=dilation,
padding=padding,
layout="NCHW",
)
)
R.output(output)
R.func_ret_value(frame.output_vars[0])
func = builder.get()
return tvm.IRModule({"main": func})
def get_avgpool_expected_codegen(input_shape, pool_size, stride, padding, pool_type, dtype):
import math
adjusted_input_shape = [
input_shape[0],
input_shape[1],
input_shape[2] + padding[0] + padding[1],
input_shape[3] + padding[2] + padding[3],
]
pool_height = math.floor(((adjusted_input_shape[2] - pool_size[0]) / stride[0]) + 1)
pool_width = math.floor(((adjusted_input_shape[3] - pool_size[1]) / stride[1]) + 1)
output_shape = [adjusted_input_shape[0], adjusted_input_shape[1], pool_height, pool_width]
attrs = {
"ceil_mode": 0,
"dilation": [1, 1],
"layout": "NCHW",
"num_inputs": 1,
"num_outputs": 1,
"out_layout": "NCHW",
"padding": list(padding),
"pool_size": pool_size,
"shape": [list(output_shape)],
"dtype": [dtype],
"strides": stride,
"count_include_pad": 0,
}
if sum(padding):
attrs["count_include_pad"] = 0
exp_codegen = [
{
"op": "input",
"name": "",
"attrs": {"shape": [list(adjusted_input_shape)], "dtype": [str(dtype)]},
},
{
"op": "kernel",
"name": "",
"inputs": [[0, 0, 0]],
"attrs": attrs,
},
]
return exp_codegen
def get_relax_reshape_mod(input_shape, output_shape, dtype):
"""
Args:
input_shape (tuple): Input tensor shape
output_shape (tuple): Desired output tensor shape
dtype (str, optional): Data type. Defaults to "float32".
Returns:
tvm.IRModule: Relax Reshape module
"""
with IRBuilder() as builder:
with relax_builder.function():
R.func_name("main")
data = R.arg("data", R.Tensor(input_shape, dtype))
with R.dataflow() as frame:
output = R.emit(R.reshape(data, output_shape))
R.output(output)
R.func_ret_value(frame.output_vars[0])
func = builder.get()
return tvm.IRModule({"main": func})
def get_relax_reshape_codegen(input_shape, output_shape, dtype):
def compute_output_shape(input_shape, output_shape):
input_elements = np.prod(input_shape)
specified_elements = np.prod([dim for dim in output_shape if dim != -1])
missing_dim = input_elements // specified_elements
return [int(dim) if dim != -1 else int(missing_dim) for dim in output_shape]
expected_output_shape = compute_output_shape(input_shape, output_shape)
expected_codegen_str = [
{
"attrs": {
"dtype": [dtype],
"shape": [list(input_shape)],
},
"name": "",
"op": "input",
},
{
"attrs": {
"dtype": [dtype],
"num_inputs": 1,
"num_outputs": 1,
"shape": [expected_output_shape],
},
"inputs": [[0, 0, 0]],
"name": "",
"op": "kernel",
},
]
return expected_codegen_str
def get_relax_global_avgpool_mod(data_shape, keepdims, dtype):
"""
Create a Relax module for Global Average Pooling (GAP).
Args:
data_shape (tuple): Input tensor shape (N, C, H, W)
dtype (str): Data type
Returns:
tvm.IRModule: Relax GAP module
"""
with IRBuilder() as builder:
with relax_builder.function():
R.func_name("main")
data = R.arg("data", R.Tensor(data_shape, dtype))
with R.dataflow() as frame:
output = R.emit(R.mean(data, axis=[2, 3], keepdims=keepdims))
R.output(output)
R.func_ret_value(frame.output_vars[0])
func = builder.get()
return tvm.IRModule({"main": func})
def get_global_avgpool_expected_codegen(input_shape, keep_dims, dtype):
"""
Generate expected codegen for Global Average Pooling.
Args:
input_shape (tuple): Input shape (N, C, H, W)
dtype (str): Data type
Returns:
dict: Expected codegen output
"""
output_shape = (
[input_shape[0], input_shape[1]]
if not keep_dims
else [input_shape[0], input_shape[1], 1, 1]
)
attrs = {
"num_inputs": 1,
"num_outputs": 1,
"shape": [list(output_shape)],
"dtype": [dtype],
"axis": [2, 3],
"keepdims": 1 if keep_dims else 0,
}
exp_codegen = [
{
"op": "input",
"name": "",
"attrs": {"shape": [list(input_shape)], "dtype": [str(dtype)]},
},
{"op": "kernel", "name": "", "inputs": [[0, 0, 0]], "attrs": attrs},
]
return exp_codegen
def get_relax_global_maxpool_mod(data_shape, keepdims, dtype):
"""
Create a Relax module for Global Average Pooling (GAP).
Args:
data_shape (tuple): Input tensor shape (N, C, H, W)
dtype (str): Data type
Returns:
tvm.IRModule: Relax GAP module
"""
N, C, H, W = data_shape
with IRBuilder() as builder:
with relax_builder.function():
R.func_name("main")
data = R.arg("data", R.Tensor(data_shape, dtype))
with R.dataflow() as frame:
output = R.emit(
R.nn.max_pool2d(
data, pool_size=(H, W), strides=(1, 1), padding=(0, 0), layout="NCHW"
)
)
R.output(output)
R.func_ret_value(frame.output_vars[0])
func = builder.get()
return tvm.IRModule({"main": func})
def get_global_maxpool_expected_codegen(input_shape, pool_size, stride, padding, pool_type, dtype):
import math
adjusted_input_shape = [
input_shape[0],
input_shape[1],
input_shape[2] + padding[0] + padding[1],
input_shape[3] + padding[2] + padding[3],
]
output_shape = [adjusted_input_shape[0], adjusted_input_shape[1], 1, 1]
attrs = {
"ceil_mode": 0,
"dilation": [1, 1],
"layout": "NCHW",
"num_inputs": 1,
"num_outputs": 1,
"out_layout": "NCHW",
"padding": padding,
"pool_size": pool_size,
"shape": [list(output_shape)],
"dtype": [dtype],
"strides": stride,
"count_include_pad": 0,
}
if sum(padding):
attrs["count_include_pad"] = 0
exp_codegen = [
{
"op": "input",
"name": "",
"attrs": {"shape": [list(adjusted_input_shape)], "dtype": [str(dtype)]},
},
{
"op": "kernel",
"name": "",
"inputs": [[0, 0, 0]],
"attrs": attrs,
},
]
return exp_codegen
def get_dequant_matmul_module(K, N):
@I.ir_module(s_tir=True)
class DequantMatmul:
@R.function
def main(
input: R.Tensor((1, "seq_len", K), dtype="float16"),
weight: R.Tensor((K // 8, N), dtype="uint32"),
scale: R.Tensor((K // 32, N), dtype="float16"),
):
seq_len = T.int64()
cls = DequantMatmul
with R.dataflow():
lv2 = relax.call_tir(
cls.dequantize,
(weight, scale),
out_ty=R.Tensor((K, N), dtype="float16"),
)
gv: R.Tensor((1, seq_len, N), dtype="float16") = relax.op.matmul(
input, lv2, out_dtype="float16"
)
R.output(gv)
return gv
@T.prim_func(s_tir=True)
def dequantize(weight: T.handle, scale: T.handle, var_dequantize: T.handle):
T.func_attr({"tirx.noalias": T.bool(True)})
lm_head_q_weight1 = T.match_buffer(weight, (T.int64(K // 8), T.int64(N)), "uint32")
lm_head_q_scale1 = T.match_buffer(scale, (T.int64(K // 32), T.int64(N)), "float16")
dequantize = T.match_buffer(var_dequantize, (T.int64(K), T.int64(N)), "float16")
# with T.sblock("root"):
compute = T.alloc_buffer((T.int64(K), T.int64(N)), "float16")
for i0, i1 in T.grid(T.int64(K), T.int64(N)):
with T.sblock("compute"):
v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
T.reads(lm_head_q_weight1[v_i0 // T.int64(8), v_i1])
T.writes(compute[v_i0, v_i1])
compute[v_i0, v_i1] = T.Cast(
"float16",
T.bitwise_and(
T.shift_right(
lm_head_q_weight1[v_i0 // T.int64(8), v_i1],
T.Cast("uint32", v_i0 % T.int64(8) * T.int64(4)),
),
T.uint32(15),
),
)
for i0, i1 in T.grid(T.int64(K), T.int64(N)):
with T.sblock("dequantize"):
v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
T.reads(compute[v_i0, v_i1], lm_head_q_scale1[v_i0 // T.int64(32), v_i1])
T.writes(dequantize[v_i0, v_i1])
dequantize[v_i0, v_i1] = (
compute[v_i0, v_i1] - T.float16(7.0)
) * lm_head_q_scale1[v_i0 // T.int64(32), v_i1]
return DequantMatmul
def get_dequant_vec_matmul_module(K, N):
@I.ir_module(s_tir=True)
class DequantVecMatmul:
@R.function
def main(
input: R.Tensor((1, 1, K), dtype="float16"),
weight: R.Tensor((K // 8, "vocab_size"), dtype="uint32"),
scale: R.Tensor((K // 32, "vocab_size"), dtype="float16"),
):
vocab_size = T.int64()
cls = DequantVecMatmul
with R.dataflow():
lv2 = relax.call_tir(
cls.dequantize,
(weight, scale),
out_ty=R.Tensor((K, vocab_size), dtype="float16"),
)
gv: R.Tensor((1, 1, vocab_size), dtype="float16") = relax.op.matmul(
input, lv2, out_dtype="float16"
)
R.output(gv)
return gv
@T.prim_func(s_tir=True)
def dequantize(weight: T.handle, scale: T.handle, var_dequantize: T.handle):
T.func_attr({"tirx.noalias": T.bool(True)})
vocab_size = T.int64()
lm_head_q_weight1 = T.match_buffer(weight, (T.int64(K // 8), vocab_size), "uint32")
lm_head_q_scale1 = T.match_buffer(scale, (T.int64(K // 32), vocab_size), "float16")
dequantize = T.match_buffer(var_dequantize, (T.int64(K), vocab_size), "float16")
# with T.sblock("root"):
compute = T.alloc_buffer((T.int64(K), vocab_size), "float16")
for i0, i1 in T.grid(T.int64(K), vocab_size):
with T.sblock("compute"):
v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
T.reads(lm_head_q_weight1[v_i0 // T.int64(8), v_i1])
T.writes(compute[v_i0, v_i1])
compute[v_i0, v_i1] = T.Cast(
"float16",
T.bitwise_and(
T.shift_right(
lm_head_q_weight1[v_i0 // T.int64(8), v_i1],
T.Cast("uint32", v_i0 % T.int64(8) * T.int64(4)),
),
T.uint32(15),
),
)
for i0, i1 in T.grid(T.int64(K), vocab_size):
with T.sblock("dequantize"):
v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
T.reads(compute[v_i0, v_i1], lm_head_q_scale1[v_i0 // T.int64(32), v_i1])
T.writes(dequantize[v_i0, v_i1])
dequantize[v_i0, v_i1] = (
compute[v_i0, v_i1] - T.float16(7.0)
) * lm_head_q_scale1[v_i0 // T.int64(32), v_i1]
return DequantVecMatmul