Skip to content

Commit 2023ce1

Browse files
chenmiao1919103yiranbin913
authored
cambricon: update and fix kernel (flagos-ai#4435)
* cambricon: update and fix kernel * cambricon: skip unsupport operators and dtype in tests * cambricon: route fused test imports through flag_gems * cambricon: align cambricon empty with the community empty.memory_format registration change --------- Co-authored-by: 103yiran <1039105206@qq.com> Co-authored-by: bin913 <842884726@qq.com>
1 parent df94a1b commit 2023ce1

277 files changed

Lines changed: 40168 additions & 1393 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎benchmark/test_blas_perf_parallel.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -203,7 +203,7 @@ def set_more_shapes(self):
203203
def get_tflops(self, op, *args, **kwargs):
204204
groups, N, K = args[1].shape
205205
size_per_group = torch.diff(
206-
args[2], prepend=torch.zeros(1, device="cuda", dtype=torch.int32)
206+
args[2], prepend=torch.zeros(1, device=args[2].device, dtype=torch.int32)
207207
)
208208
total_flops = 0
209209
for i in range(groups):
@@ -366,8 +366,8 @@ def group_mm_input_fn(groups, N, K, cur_dtype, device):
366366
M_g = random.randint(1, 16384)
367367
N_g = N
368368
K_g = K
369-
A_g = torch.rand([M_g, K_g], device="cuda", dtype=cur_dtype)
370-
B_g = torch.rand([K_g, N_g], device="cuda", dtype=cur_dtype)
369+
A_g = torch.rand([M_g, K_g], device=device, dtype=cur_dtype)
370+
B_g = torch.rand([K_g, N_g], device=device, dtype=cur_dtype)
371371
group_A_list.append(A_g)
372372
group_B_list.append(B_g)
373373
M_list.append(M_g)
@@ -377,7 +377,7 @@ def group_mm_input_fn(groups, N, K, cur_dtype, device):
377377
mat_a = torch.cat([x for x in group_A_list], dim=0)
378378
mat_b = torch.stack([x for x in group_B_list], dim=0)
379379
offs = torch.tensor(
380-
[sum(M_list[: i + 1]) for i in range(groups)], dtype=torch.int32, device="cuda"
380+
[sum(M_list[: i + 1]) for i in range(groups)], dtype=torch.int32, device=device
381381
)
382382

383383
yield mat_a, mat_b, offs

‎benchmark/test_fused_deepseek_v4_qnorm_rope_kv_rope_insert.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717
import pytest
1818
import torch
1919

20-
from flag_gems.fused import fused_deepseek_v4_qnorm_rope_kv_rope_insert
20+
from flag_gems import fused_deepseek_v4_qnorm_rope_kv_rope_insert
2121

2222
from . import base
2323

‎benchmark/test_fused_moe_int4_w4a16.py‎

Lines changed: 31 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -87,20 +87,37 @@ def _int4_w4a16_input_fn(self, config, dtype):
8787
dtype=torch.int8,
8888
)
8989
for e in range(num_experts):
90-
w1_int4[e] = torch.randint(
91-
-8,
92-
8,
93-
(intermediate_size * 2, hidden_size),
94-
device=device,
95-
dtype=torch.int8,
96-
)
97-
w2_int4[e] = torch.randint(
98-
-8,
99-
8,
100-
(hidden_size, intermediate_size),
101-
device=device,
102-
dtype=torch.int8,
103-
)
90+
if flag_gems.vendor_name == "cambricon":
91+
# Cambricon torch.randint currently does not support int8/int16 generation.
92+
w1_int4[e] = torch.randint(
93+
-8,
94+
8,
95+
(intermediate_size * 2, hidden_size),
96+
device="cpu",
97+
dtype=torch.int8,
98+
).to(device)
99+
w2_int4[e] = torch.randint(
100+
-8,
101+
8,
102+
(hidden_size, intermediate_size),
103+
device="cpu",
104+
dtype=torch.int8,
105+
).to(device)
106+
else:
107+
w1_int4[e] = torch.randint(
108+
-8,
109+
8,
110+
(intermediate_size * 2, hidden_size),
111+
device=device,
112+
dtype=torch.int8,
113+
)
114+
w2_int4[e] = torch.randint(
115+
-8,
116+
8,
117+
(hidden_size, intermediate_size),
118+
device=device,
119+
dtype=torch.int8,
120+
)
104121

105122
# Per-channel scales [E, output_dim]
106123
w1_scale = (

‎benchmark/test_gcd.py‎

Lines changed: 33 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,8 @@
1515
import pytest
1616
import torch
1717

18+
import flag_gems
19+
1820
from . import base, consts
1921

2022

@@ -29,20 +31,37 @@ def test_gcd():
2931

3032

3133
def gcd_out_input_fn(shape, dtype, device):
32-
inp1 = torch.randint(
33-
torch.iinfo(dtype).min,
34-
torch.iinfo(dtype).max,
35-
shape,
36-
dtype=dtype,
37-
device=device,
38-
)
39-
inp2 = torch.randint(
40-
torch.iinfo(dtype).min,
41-
torch.iinfo(dtype).max,
42-
shape,
43-
dtype=dtype,
44-
device=device,
45-
)
34+
if flag_gems.vendor_name == "cambricon":
35+
# Cambricon torch.randint currently does not support int8/int16 generation.
36+
inp1 = torch.randint(
37+
torch.iinfo(dtype).min,
38+
torch.iinfo(dtype).max,
39+
shape,
40+
dtype=dtype,
41+
device="cpu",
42+
).to(device)
43+
inp2 = torch.randint(
44+
torch.iinfo(dtype).min,
45+
torch.iinfo(dtype).max,
46+
shape,
47+
dtype=dtype,
48+
device="cpu",
49+
).to(device)
50+
else:
51+
inp1 = torch.randint(
52+
torch.iinfo(dtype).min,
53+
torch.iinfo(dtype).max,
54+
shape,
55+
dtype=dtype,
56+
device=device,
57+
)
58+
inp2 = torch.randint(
59+
torch.iinfo(dtype).min,
60+
torch.iinfo(dtype).max,
61+
shape,
62+
dtype=dtype,
63+
device=device,
64+
)
4665
out = torch.empty(shape, dtype=dtype, device=device)
4766
yield inp1, inp2, {"out": out}
4867

‎benchmark/test_grouped_mm.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,7 @@ def set_more_shapes(self):
3434
def get_tflops(self, op, *args, **kwargs):
3535
groups, N, K = args[1].shape
3636
size_per_group = torch.diff(
37-
args[2], prepend=torch.zeros(1, device="cuda", dtype=torch.int32)
37+
args[2], prepend=torch.zeros(1, device=args[2].device, dtype=torch.int32)
3838
)
3939
total_flops = 0
4040
for i in range(groups):
@@ -54,8 +54,8 @@ def _input_fn(groups, N, K, cur_dtype, device):
5454
M_g = random.randint(1, 16384)
5555
N_g = N
5656
K_g = K
57-
A_g = torch.rand([M_g, K_g], device="cuda", dtype=cur_dtype)
58-
B_g = torch.rand([K_g, N_g], device="cuda", dtype=cur_dtype)
57+
A_g = torch.rand([M_g, K_g], device=device, dtype=cur_dtype)
58+
B_g = torch.rand([K_g, N_g], device=device, dtype=cur_dtype)
5959
group_A_list.append(A_g)
6060
group_B_list.append(B_g)
6161
M_list.append(M_g)
@@ -65,7 +65,7 @@ def _input_fn(groups, N, K, cur_dtype, device):
6565
mat_a = torch.cat([x for x in group_A_list], dim=0)
6666
mat_b = torch.stack([x for x in group_B_list], dim=0)
6767
offs = torch.tensor(
68-
[sum(M_list[: i + 1]) for i in range(groups)], dtype=torch.int32, device="cuda"
68+
[sum(M_list[: i + 1]) for i in range(groups)], dtype=torch.int32, device=device
6969
)
7070

7171
yield mat_a, mat_b, offs

‎benchmark/test_margin_ranking_loss.py‎

Lines changed: 24 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -15,24 +15,42 @@
1515
import pytest
1616
import torch
1717

18+
import flag_gems
19+
1820
from . import base, consts
1921

2022

2123
def _input_fn(shape, dtype, device):
2224
inp1 = torch.randn(shape, dtype=dtype, device=device)
2325
inp2 = torch.randn(shape, dtype=dtype, device=device)
24-
target = (torch.randint(0, 2, shape, device=device, dtype=torch.int8) * 2 - 1).to(
25-
dtype
26-
)
26+
if flag_gems.vendor_name == "cambricon":
27+
# Cambricon torch.randint currently does not support int8/int16 generation.
28+
target = (
29+
(torch.randint(0, 2, shape, device="cpu", dtype=torch.int8) * 2 - 1)
30+
.to(dtype)
31+
.to(device)
32+
)
33+
else:
34+
target = (
35+
torch.randint(0, 2, shape, device=device, dtype=torch.int8) * 2 - 1
36+
).to(dtype)
2737
yield inp1, inp2, target, 0.5, 1
2838

2939

3040
def _backward_input_fn(shape, dtype, device):
3141
inp1 = torch.randn(shape, dtype=dtype, device=device, requires_grad=True)
3242
inp2 = torch.randn(shape, dtype=dtype, device=device, requires_grad=True)
33-
target = (torch.randint(0, 2, shape, device=device, dtype=torch.int8) * 2 - 1).to(
34-
dtype
35-
)
43+
if flag_gems.vendor_name == "cambricon":
44+
# Cambricon torch.randint currently does not support int8/int16 generation.
45+
target = (
46+
(torch.randint(0, 2, shape, device="cpu", dtype=torch.int8) * 2 - 1)
47+
.to(dtype)
48+
.to(device)
49+
)
50+
else:
51+
target = (
52+
torch.randint(0, 2, shape, device=device, dtype=torch.int8) * 2 - 1
53+
).to(dtype)
3654
yield inp1, inp2, target, 0.5, 1
3755

3856

‎benchmark/test_masked_scale.py‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,8 @@
1717
import pytest
1818
import torch
1919

20+
import flag_gems
21+
2022
from . import base, utils
2123

2224

@@ -29,7 +31,13 @@ def set_more_shapes(self):
2931
def get_input_iter(self, cur_dtype) -> Generator:
3032
for shape in self.shapes:
3133
inp = utils.generate_tensor_input(shape, cur_dtype, self.device)
32-
mask = torch.randint(0, 2, shape, dtype=torch.uint8, device=self.device)
34+
if flag_gems.vendor_name == "cambricon":
35+
# Cambricon torch.randint currently does not support uint8 generation.
36+
mask = torch.randint(0, 2, shape, dtype=torch.uint8, device="cpu").to(
37+
self.device
38+
)
39+
else:
40+
mask = torch.randint(0, 2, shape, dtype=torch.uint8, device=self.device)
3341
scale = 2.0
3442
yield inp, mask, scale
3543

‎benchmark/test_nested_view_from_buffer_copy.py‎

Lines changed: 15 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -2,24 +2,27 @@
22

33
# The concurrent agent runner may leave a stale editable install redirector pointing to
44
# a different worktree. Remove it so our worktree's flag_gems is loaded.
5-
sys.meta_path = [
6-
m for m in sys.meta_path if "ScikitBuildRedirectingFinder" not in str(type(m))
7-
]
5+
import flag_gems
86

9-
# Swap our worktree's src to the front of sys.path.
10-
_our_src = "/tmp/flaggems_agent_worktrees/agent__nested_view_from_buffer_copy_179605-1781252494/src"
11-
sys.path = [_our_src] + [p for p in sys.path if p != _our_src]
7+
if flag_gems.vendor_name != "cambricon":
8+
sys.meta_path = [
9+
m for m in sys.meta_path if "ScikitBuildRedirectingFinder" not in str(type(m))
10+
]
1211

13-
# Clear any cached flag_gems modules.
14-
for _k in list(sys.modules):
15-
if "flag_gems" in _k:
16-
del sys.modules[_k]
12+
# Swap our worktree's src to the front of sys.path.
13+
_our_src = "/tmp/flaggems_agent_worktrees/agent__nested_view_from_buffer_copy_179605-1781252494/src"
14+
sys.path = [_our_src] + [p for p in sys.path if p != _our_src]
15+
16+
# Clear any cached flag_gems modules.
17+
for _k in list(sys.modules):
18+
if "flag_gems" in _k:
19+
del sys.modules[_k]
20+
21+
import flag_gems # noqa: E402
1722

1823
import pytest # noqa: E402
1924
import torch # noqa: E402
2025

21-
import flag_gems # noqa: E402
22-
2326
from . import base, consts # noqa: E402
2427

2528

‎benchmark/test_remainder.py‎

Lines changed: 36 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -62,20 +62,37 @@ def test_remainder_tensor_inplace():
6262

6363

6464
def remainder_scalar_input_fn(shape, dtype, device):
65-
inp = torch.randint(
66-
torch.iinfo(dtype).min,
67-
torch.iinfo(dtype).max,
68-
shape,
69-
dtype=dtype,
70-
device=device,
71-
)
72-
scalar = torch.randint(
73-
torch.iinfo(dtype).min,
74-
torch.iinfo(dtype).max,
75-
(1,),
76-
dtype=dtype,
77-
device=device,
78-
).item()
65+
if flag_gems.vendor_name == "cambricon":
66+
# Cambricon torch.randint currently does not support int8/int16 generation.
67+
inp = torch.randint(
68+
torch.iinfo(dtype).min,
69+
torch.iinfo(dtype).max,
70+
shape,
71+
dtype=dtype,
72+
device="cpu",
73+
).to(device)
74+
scalar = torch.randint(
75+
torch.iinfo(dtype).min,
76+
torch.iinfo(dtype).max,
77+
(1,),
78+
dtype=dtype,
79+
device="cpu",
80+
).item()
81+
else:
82+
inp = torch.randint(
83+
torch.iinfo(dtype).min,
84+
torch.iinfo(dtype).max,
85+
shape,
86+
dtype=dtype,
87+
device=device,
88+
)
89+
scalar = torch.randint(
90+
torch.iinfo(dtype).min,
91+
torch.iinfo(dtype).max,
92+
(1,),
93+
dtype=dtype,
94+
device=device,
95+
).item()
7996
if scalar == 0:
8097
scalar = 1
8198
yield inp, scalar
@@ -111,7 +128,11 @@ def test_remainder_scalar_inplace():
111128

112129

113130
def scalar_tensor_remainder_input_fn(shape, dtype, device):
114-
inp = torch.randint(1, 100, shape, dtype=dtype, device=device)
131+
if flag_gems.vendor_name == "cambricon":
132+
# Cambricon torch.randint currently does not support int8/int16 generation.
133+
inp = torch.randint(1, 100, shape, dtype=dtype, device="cpu").to(device)
134+
else:
135+
inp = torch.randint(1, 100, shape, dtype=dtype, device=device)
115136
scalar = 7
116137
yield scalar, inp
117138

‎benchmark/test_router_gemm.py‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,14 @@ def torch_router_gemm(x, weight):
4545
router_gemm = None
4646
ROUTER_GEMM_AVAILABLE = False
4747

48+
if base.vendor_name == "cambricon":
49+
try:
50+
from flag_gems.runtime.backend._cambricon.ops.mm import router_gemm
51+
52+
ROUTER_GEMM_AVAILABLE = True
53+
except Exception:
54+
pass
55+
4856

4957
class RouterGemmBenchmark(base.Benchmark):
5058
DEFAULT_METRICS = consts.DEFAULT_METRICS[:] + ["tflops"]
@@ -73,7 +81,7 @@ def get_tflops(self, op, *args, **kwargs):
7381
@pytest.mark.router_gemm
7482
@pytest.mark.skipif(
7583
not ROUTER_GEMM_AVAILABLE,
76-
reason="router_gemm benchmark requires NVIDIA Hopper backend",
84+
reason="router_gemm benchmark requires Cambricon or NVIDIA Hopper backend",
7785
)
7886
def test_perf_router_gemm():
7987
bench = RouterGemmBenchmark(

0 commit comments

Comments
 (0)