GMMFR adds support for cumsum in groulist and none option in shareInput.-op-plugin-AtomGit

Mop_plugin/ops/opapi/GroupedMatmulFinalizeRoutingKernelNpuOpApi.cpp+3-24

	@@ -48,8 +48,8 @@ at::Tensor npu_grouped_matmul_finalize_routing(
48	)	48	)
49	{	49	{
50	bool is_weight_nz = is_nz_format(w);	50	bool is_weight_nz = is_nz_format(w);
51	- TORCH_CHECK(group_list_type == 1,	51	+ TORCH_CHECK(group_list_type == 1 \|\| group_list_type == 0,
52	- "only support group_list_type's value is 1.",	52	+ "only support group_list_type's value 0 or 1.",
53	OPS_ERROR(ErrCode::PARAM));	53	OPS_ERROR(ErrCode::PARAM));
54		54
55	auto x_dim_num = x.dim();	55	auto x_dim_num = x.dim();
	@@ -69,11 +69,6 @@ at::Tensor npu_grouped_matmul_finalize_routing(
69		69
70	auto output_size = op_infer::array_to_small_vector(x.sizes());	70	auto output_size = op_infer::array_to_small_vector(x.sizes());
71	int32_t output_bs_real = static_cast<int32_t>(output_bs.value_or(0));	71	int32_t output_bs_real = static_cast<int32_t>(output_bs.value_or(0));
72	- if (!shared_input.has_value() && !logit.has_value()) {
73	- TORCH_CHECK(output_bs_real == x_m_dim,
74	- "When shared_input and logit is None, output_bs must equal to M",
75	- OPS_ERROR(ErrCode::PARAM));
76	- }
77	if (output_bs_real == 0) {	72	if (output_bs_real == 0) {
78	output_bs_real = x_m_dim;	73	output_bs_real = x_m_dim;
79	}	74	}
	@@ -97,28 +92,12 @@ at::Tensor npu_grouped_matmul_finalize_routing(
97	auto tuning_config_real = tuning_config.value_or(at::IntArrayRef{});	92	auto tuning_config_real = tuning_config.value_or(at::IntArrayRef{});
98	auto antiquant_scale_real = at::Tensor();	93	auto antiquant_scale_real = at::Tensor();
99	auto antiquant_offset_real = at::Tensor();	94	auto antiquant_offset_real = at::Tensor();
100	-
101	- auto scene_has_share = false;
102	- auto scene_no_share = false;
103	- if (scale.has_value() && shared_input.has_value()
104	- && logit.has_value() && row_index.has_value()) {
105	- scene_has_share = true;
106	- }
107	-
108	- if (scale.has_value() && !shared_input.has_value()
109	- && !logit.has_value() && row_index.has_value()) {
110	- scene_no_share = true;
111	- }
112	-
113	- TORCH_CHECK(scene_has_share \|\| scene_no_share,
114	- "input tensor only support shared_input and logit empty tensor",
115	- OPS_ERROR(ErrCode::PARAM));
116		95
117	at::ScalarType dst_type = c10::value_or_else(dtype, [] {return at::ScalarType::Float;});	96	at::ScalarType dst_type = c10::value_or_else(dtype, [] {return at::ScalarType::Float;});
118	TORCH_CHECK(dst_type == at::ScalarType::Float,	97	TORCH_CHECK(dst_type == at::ScalarType::Float,
119	"The dtype should be float", OPS_ERROR(ErrCode::PARAM));	98	"The dtype should be float", OPS_ERROR(ErrCode::PARAM));
120		99
121	- if (shared_input.has_value() && logit.has_value()) {	100	+ if (shared_input.has_value()) {
122	TORCH_CHECK(dst_type == at::ScalarType::Float,	101	TORCH_CHECK(dst_type == at::ScalarType::Float,
123	"When shared_input and logit is not None, the dtype must be float32",	102	"When shared_input and logit is not None, the dtype must be float32",
124	OPS_ERROR(ErrCode::PARAM));	103	OPS_ERROR(ErrCode::PARAM));

Mop_plugin/python/meta/_meta_registrations.py+1-13

	@@ -1759,25 +1759,13 @@ def npu_grouped_matmul_finalize_routing_meta(x, w, group_list, *, scale=None, bi
1759		1759
1760	torch._check(	1760	torch._check(
1761	x_dim == 2 and w_dim == 3,	1761	x_dim == 2 and w_dim == 3,
1762	- lambda: "input tensor only support shared_input and logit empty tensor" + ops_error(ErrCode.VALUE),	1762	+ lambda: "x_dim should be 2 and w_dim should be 3." + ops_error(ErrCode.VALUE),
1763	)	1763	)
1764	torch._check(	1764	torch._check(
1765	dimn > 0,	1765	dimn > 0,
1766	lambda: "n value must bigger than 0." + ops_error(ErrCode.VALUE),	1766	lambda: "n value must bigger than 0." + ops_error(ErrCode.VALUE),
1767	)	1767	)
1768		1768
1769	- scene1 = False
1770	- scene2 = False
1771	- scene1 = (scale is not None and
1772	- group_list is not None and shared_input is not None and
1773	- logit is not None and row_index is not None)
1774	- scene2 = (scale is not None and
1775	- group_list is not None and shared_input is None and
1776	- logit is None and row_index is not None)
1777	- torch._check(
1778	- scene1 or scene2,
1779	- lambda: "input tensor only support shared_input and logit empty tensor" + ops_error(ErrCode.VALUE),
1780	- )
1781	if dtype is None:	1769	if dtype is None:
1782	dtype = torch.float32	1770	dtype = torch.float32
1783	if shared_input is not None and logit is not None:	1771	if shared_input is not None and logit is not None:

Mtest/core_tests/test_fake_tensor.py+33-0

	@@ -2464,6 +2464,39 @@ class TestGMMFinalizeRouting(TestCase):
2464	expect_ret = torch.normal(0, 0.1, (output_bs, n), dtype=torch.float32)	2464	expect_ret = torch.normal(0, 0.1, (output_bs, n), dtype=torch.float32)
2465	self.assertTrue(result.shape == expect_ret.shape)	2465	self.assertTrue(result.shape == expect_ret.shape)
2466	self.assertTrue(result.dtype == expect_ret.dtype)	2466	self.assertTrue(result.dtype == expect_ret.dtype)
		2467	+
		2468	+ @unittest.skip("Skipping due to outdated CANN version; please update CANN to the latest version and remove this skip.")
		2469	+ def test_npu_grouped_matmul_finalise_routing_sharedinput_none_grouplist_cumsum_meta(self):
		2470	+ with FakeTensorMode():
		2471	+ m, k, n, batch, topK, group_num, shared_input_scale = 576, 2048, 7168, 72, 8, 8, 1
		2472	+ x = torch.randint(-10, 10, (m, k), dtype=torch.int8)
		2473	+ weight = torch.randint(-10, 10, (group_num, k, n), dtype=torch.int8)
		2474	+ scale = torch.normal(0, 0.01, (group_num, n), dtype=torch.float32)
		2475	+ pertoken_scale = torch.normal(0, 0.01, (m, 1), dtype=torch.float32)
		2476	+ group_list = torch.tensor([batch] * group_num, dtype=torch.int64)
		2477	+ logit_ori = torch.normal(0, 0.1, (batch, group_num), dtype=torch.float32)
		2478	+ routing = torch.argsort(logit_ori, 1)[:, -topK:]
		2479	+ logit = F.softmax(
		2480	+ logit_ori[torch.arange(batch).reshape(-1, 1).repeat(1, topK), routing],
		2481	+ dim=1,
		2482	+ dtype=torch.float32
		2483	+ ).reshape(m)
		2484	+ row_index = (torch.argsort(routing.reshape(-1)) // topK).to(torch.int64)
		2485	+ shared_input_offset = batch // 2
		2486	+ output_bs = batch
		2487	+ group_list_type = 0
		2488	+ group_list = torch.cumsum(group_list, dim=0)
		2489	+ weightNz = torch_npu.npu_format_cast(weight.npu(), 29)
		2490	+ pertoken_scale = pertoken_scale.reshape(m)
		2491	+ result = torch_npu.npu_grouped_matmul_finalize_routing(
		2492	+ x.npu(), weightNz, group_list.npu(), scale=scale.npu(),
		2493	+ pertoken_scale=pertoken_scale.npu(), shared_input=None,
		2494	+ logit=logit.npu(), row_index=row_index.npu(),
		2495	+ shared_input_offset=shared_input_offset, output_bs=output_bs, group_list_type=group_list_type
		2496	+ ).to("cpu")
		2497	+ expect_ret = torch.normal(0, 0.1, (output_bs, n), dtype=torch.float32)
		2498	+ self.assertTrue(result.shape == expect_ret.shape)
		2499	+ self.assertTrue(result.dtype == expect_ret.dtype)
2467		2500
2468		2501
2469	class TestTransposeBatchMatmul(TestCase):	2502	class TestTransposeBatchMatmul(TestCase):

Mtest/test_custom_ops/test_npu_grouped_matmul_finalize_routing.py+21-34

	@@ -74,8 +74,9 @@ class TestGroupedMatmulFinalizeRouting(TestCase):
74	out[shared_input_offset:shared_input_offset + shared_input.shape[0], :] += \	74	out[shared_input_offset:shared_input_offset + shared_input.shape[0], :] += \
75	shared_input_scale * shared_input.to(torch.float32)	75	shared_input_scale * shared_input.to(torch.float32)
76	else:	76	else:
		77	+ out = mm_out * logit.reshape(-1, 1)
77	index = torch.argsort(row_index, dim=0)	78	index = torch.argsort(row_index, dim=0)
78	- out = mm_out[index].reshape(-1, topK, mm_out.shape[-1]).sum(dim=1)	79	+ out = out[index].reshape(-1, topK, mm_out.shape[-1]).sum(dim=1)
79	return out	80	return out
80		81
81	def supported_a8w4_op_exec(self, topK, x_in, weight_in, groupList_in, scale_in,	82	def supported_a8w4_op_exec(self, topK, x_in, weight_in, groupList_in, scale_in,
	@@ -135,53 +136,39 @@ class TestGroupedMatmulFinalizeRouting(TestCase):
135	).to("cpu")	136	).to("cpu")
136	self.assertRtolEqual(supported_output, custom_output, 0.001)	137	self.assertRtolEqual(supported_output, custom_output, 0.001)
137		138
		139	+ @unittest.skip("Skipping due to outdated CANN version; please update CANN to the latest version and remove this skip.")
138	@SupportedDevices(["Ascend910B"])	140	@SupportedDevices(["Ascend910B"])
139	- def test_npu_grouped_matmul_finalize_routing_2(self, device="npu"):	141	+ def test_npu_grouped_matmul_finalize_routing_sharedinput_none_grouplist_cumsum(self, device="npu"):
140	- m, k, n, batch, topK, group_num = 72, 2048, 7168, 72, 1, 1	142	+ m, k, n, batch, topK, group_num, shared_input_scale = 576, 2048, 7168, 72, 8, 8, 1
141	x = torch.randint(-10, 10, (m, k), dtype=torch.int8)	143	x = torch.randint(-10, 10, (m, k), dtype=torch.int8)
142	weight = torch.randint(-10, 10, (group_num, k, n), dtype=torch.int8)	144	weight = torch.randint(-10, 10, (group_num, k, n), dtype=torch.int8)
143	scale = torch.normal(0, 0.01, (group_num, n), dtype=torch.float32)	145	scale = torch.normal(0, 0.01, (group_num, n), dtype=torch.float32)
144	pertoken_scale = torch.normal(0, 0.01, (m, 1), dtype=torch.float32)	146	pertoken_scale = torch.normal(0, 0.01, (m, 1), dtype=torch.float32)
145	group_list = torch.tensor([batch] * group_num, dtype=torch.int64)	147	group_list = torch.tensor([batch] * group_num, dtype=torch.int64)
146	-
147	logit_ori = torch.normal(0, 0.1, (batch, group_num), dtype=torch.float32)	148	logit_ori = torch.normal(0, 0.1, (batch, group_num), dtype=torch.float32)
148	routing = torch.argsort(logit_ori, 1)[:, -topK:]	149	routing = torch.argsort(logit_ori, 1)[:, -topK:]
		150	+ logit = F.softmax(
		151	+ logit_ori[torch.arange(batch).reshape(-1, 1).repeat(1, topK), routing],
		152	+ dim=1,
		153	+ dtype=torch.float32
		154	+ ).reshape(m)
149	row_index = (torch.argsort(routing.reshape(-1)) // topK).to(torch.int64)	155	row_index = (torch.argsort(routing.reshape(-1)) // topK).to(torch.int64)
150	- output_bs = m	156	+ shared_input_offset = batch // 2
		157	+ output_bs = batch
151		158
152	- supported_output = self.supported_op_exec(topK, x, weight, group_list,	159	+ supported_output = self.supported_op_exec(topK, x, weight, group_list, scale,
153	- scale, pertoken_scale, row_index=row_index)	160	+ pertoken_scale, logit=logit, row_index=row_index,
		161	+ shared_input_scale=shared_input_scale,
		162	+ shared_input_offset=shared_input_offset)
		163	+ group_list_type = 0
		164	+ group_list = torch.cumsum(group_list, dim=0)
154	weightNz = torch_npu.npu_format_cast(weight.npu(), 29)	165	weightNz = torch_npu.npu_format_cast(weight.npu(), 29)
155	pertoken_scale = pertoken_scale.reshape(m)	166	pertoken_scale = pertoken_scale.reshape(m)
156	custom_output = torch_npu.npu_grouped_matmul_finalize_routing(	167	custom_output = torch_npu.npu_grouped_matmul_finalize_routing(
157	x.npu(), weightNz, group_list.npu(), scale=scale.npu(),	168	x.npu(), weightNz, group_list.npu(), scale=scale.npu(),
158	- pertoken_scale=pertoken_scale.npu(), row_index=row_index.npu(),	169	+ pertoken_scale=pertoken_scale.npu(), shared_input=None,
159	- output_bs=output_bs	170	+ logit=logit.npu(), row_index=row_index.npu(),
160	- ).to("cpu")	171	+ shared_input_offset=shared_input_offset, output_bs=output_bs, group_list_type=group_list_type
161	- self.assertRtolEqual(supported_output, custom_output, 0.001)
162	-
163	- @SupportedDevices(["Ascend910B"])
164	- def test_npu_grouped_matmul_finalize_routing_3(self, device="npu"):
165	- m, k, n, batch, topK, group_num = 72, 2048, 7168, 72, 1, 1
166	- x = torch.randint(-10, 10, (m, k), dtype=torch.int8)
167	- weight = torch.randint(-10, 10, (group_num, k, n), dtype=torch.int8)
168	- scale = torch.normal(0, 0.01, (group_num, n), dtype=torch.float32)
169	- pertoken_scale = torch.normal(0, 0.01, (m, 1), dtype=torch.float32)
170	- group_list = torch.tensor([batch] * group_num, dtype=torch.int64)
171	-
172	- logit_ori = torch.normal(0, 0.1, (batch, group_num), dtype=torch.float32)
173	- routing = torch.argsort(logit_ori, 1)[:, -topK:]
174	- row_index = (torch.argsort(routing.reshape(-1)) // topK).to(torch.int64)
175	- output_bs = m
176	-
177	- supported_output = self.supported_op_exec(topK, x, weight, group_list,
178	- scale, pertoken_scale, row_index=row_index)
179	- weightNz = torch_npu.npu_format_cast(weight.npu(), 29)
180	- pertoken_scale = pertoken_scale.reshape(m)
181	- custom_output = torch_npu.npu_grouped_matmul_finalize_routing(
182	- x.npu(), weightNz, group_list.npu(), scale=scale.npu(),
183	- pertoken_scale=pertoken_scale.npu(), row_index=row_index.npu(),
184	- output_bs=output_bs, tuning_config=[8]
185	).to("cpu")	172	).to("cpu")
186	self.assertRtolEqual(supported_output, custom_output, 0.001)	173	self.assertRtolEqual(supported_output, custom_output, 0.001)
187		174