[A5 TQuant]: Adding support for 2D DN MXFP4 Quantization-pto-isa-AtomGit

Minclude/pto/common/pto_instr.hpp+12-0

	@@ -2327,6 +2327,18 @@ PTO_INST RecordEvent TQUANT(TileDataOut &dst, TileDataSrc &src, TileDataExp *exp
2327	return {};	2327	return {};
2328	}	2328	}
2329		2329
		2330	+template <int grp_axis, auto mx_alg, typename TileDataOut = void, typename TileDataSrc = void,
		2331	+ typename TileDataExp = void, typename TileDataMax = void, typename TileDataScaling = void,
		2332	+ typename... WaitEvents>
		2333	+PTO_INST RecordEvent TQuant(TileDataOut &dst, TileDataSrc &src, TileDataExp exp, TileDataMax max,
		2334	+ TileDataScaling *scaling, WaitEvents &...events)
		2335	+{
		2336	+ TSYNC(events...);
		2337	+ TQUANT_IMPL<grp_axis, mx_alg, TileDataOut, TileDataSrc, TileDataExp, TileDataMax, TileDataScaling>(dst, src, exp,
		2338	+ max, scaling);
		2339	+ return {};
		2340	+}
		2341	+
2330	template <typename TileDataDst, typename TileDataSrc, typename... WaitEvents>	2342	template <typename TileDataDst, typename TileDataSrc, typename... WaitEvents>
2331	PTO_INST RecordEvent TINTERLEAVE(TileDataDst &dst1, TileDataDst &dst0, TileDataSrc &src1, TileDataSrc &src0,	2343	PTO_INST RecordEvent TINTERLEAVE(TileDataDst &dst1, TileDataDst &dst0, TileDataSrc &src1, TileDataSrc &src0,
2332	WaitEvents &...events)	2344	WaitEvents &...events)

Minclude/pto/common/type.hpp+11-0

	@@ -215,6 +215,17 @@ enum class QuantScaleAlg
215	NV = 1,	215	NV = 1,
216	};	216	};
217		217
		218	+// Unified MX-quant algorithm tag: encodes MX destination format
		219	+// (FP8_E4M3 / FP4_E2M1) and the shared-exponent scale algorithm (OCP / NV) in a
		220	+// single enum
		221	+enum class MxQuantAlg
		222	+{
		223	+ OcpMxFp8E4M3 = 0, // QuantType::MXFP8 + QuantScaleAlg::OCP
		224	+ NvMxFp8E4M3 = 1, // QuantType::MXFP8 + QuantScaleAlg::NV
		225	+ OcpMxFp4E2M1 = 2, // QuantType::MXFP4_E2M1 + QuantScaleAlg::OCP
		226	+ NvMxFp4E2M1 = 3, // QuantType::MXFP4_E2M1 + QuantScaleAlg::NV
		227	+};
		228	+
218	enum class VecStoreMode	229	enum class VecStoreMode
219	{	230	{
220	ND = 0,	231	ND = 0,

Minclude/pto/npu/a5/TQuant.hpp+1072-655

Minclude/pto/npu/a5/datatype.hpp+2-2

	@@ -14,7 +14,7 @@ See LICENSE in the root of the software repository for the full text of the Lice
14	namespace pto {	14	namespace pto {
15	template <typename T>	15	template <typename T>
16	struct TypeGet;	16	struct TypeGet;
17	-#if defined(__DAV_VEC__)	17	+#if defined(__DAV_VEC__) \|\| (defined(__CCE_AICORE__) && !defined(__CPU_SIM) && !defined(__COSTMODEL))
18	template <>	18	template <>
19	struct TypeGet<bfloat16_t> {	19	struct TypeGet<bfloat16_t> {
20	using T = vector_bf16;	20	using T = vector_bf16;
	@@ -126,7 +126,7 @@ template <>
126	struct TypeGet<vector_s8> {	126	struct TypeGet<vector_s8> {
127	using T = vector_s8;	127	using T = vector_s8;
128	};	128	};
129	-#if defined(__DAV_VEC__)	129	+#if defined(__DAV_VEC__) \|\| (defined(__CCE_AICORE__) && !defined(__CPU_SIM) && !defined(__COSTMODEL))
130	template <>	130	template <>
131	struct TypeGet<vector_bf16> {	131	struct TypeGet<vector_bf16> {
132	using T = vector_bf16;	132	using T = vector_bf16;

Mtests/npu/a5/src/st/testcase/tquant_dn/gen_data.py+211-0

@@ -14,6 +14,7 @@ import math

14 import os 14 import os

15 15

16 import numpy as np 16 import numpy as np

17 +from ml_dtypes import bfloat16, float4_e2m1fn

17 18

18 19

19 def fp32_to_bf16_bits(x): 20 def fp32_to_bf16_bits(x):

@@ -132,6 +133,138 @@ def quant_bf16_to_mxfp8_dn(src_bf16_fp32, m, n_pad):

132 return fp8_nd, e8_dn, fp8_nz, e8_zz 133 return fp8_nd, e8_dn, fp8_nz, e8_zz

133 134

134 135

136 +# ---------------------------------------------------------------------------

137 +# MXFP4 (E2M1) DN stage. Group max is taken across 32 rows per column (DN

138 +# layout: groups live on axis 0), producing [hat_m, n] max/e8/scaling tiles.

139 +#

140 +# The two source paths each mirror the *exact* hardware data flow so the

141 +# golden is an independent verification (via ml_dtypes), not a port of CCE:

142 +# * bf16 source : bf16 * bf16 (bf16 rounding) -> NaN->+Inf -> bf16->fp4

143 +# * fp16 source : fp16->f32, f32*scale(f32), f32->bf16(round) -> bf16->fp4

144 +#

145 +# OCP MXFP4 shared exponent is clamped at max_exp = 0x0100 (E2M1 max code),

146 +# derived from the bf16 exponent of the per-group abs-max (see OcpMxFp4E2M1Spec).

147 +# ---------------------------------------------------------------------------

148 +def bf16_maxes_to_e2m1_dn(group_max_bf16_fp32):

149 + """OCP MXFP4 e8m0 + bf16 scaling from per-group abs max (fp32, exact bf16 exp)."""

150 + bits = np.asarray(group_max_bf16_fp32, dtype=np.float32).view(np.uint32)

151 + u16 = (bits >> 16).astype(np.uint16) # bf16 rounding view

152 + exp_bf16 = u16 & np.uint16(0x7F80)

153 + mant_bf16 = u16 & np.uint16(0x007F)

154 +

155 + nan_mask = (exp_bf16 == 0x7F80) & (mant_bf16 != 0)

156 + # Clamp exponent to E2M1 maximum (0x0100 -> shared_exp baseline).

157 + exp_clamped = np.maximum(exp_bf16, np.uint16(0x0100))

158 + shared_exp_bits = exp_clamped.astype(np.uint16) - np.uint16(0x0100)

159 + e8m0 = ((shared_exp_bits >> 7) & np.uint16(0xFF)).astype(np.uint8)

160 + scale_bits = np.uint16(0x7F00) - shared_exp_bits

161 + scale_bits = np.where(nan_mask, np.uint16(0x7FC0), scale_bits) # NaN scaling

162 + e8m0 = np.where(nan_mask, np.uint8(0xFF), e8m0)

163 +

164 + scaling_fp32 = (scale_bits.astype(np.uint32) << 16).view(np.float32)

165 + return e8m0, scaling_fp32, scale_bits

166 +

167 +

168 +def _bf16_mul_round(x_bf16_fp32, scale_bf16_fp32):

169 + """Hardware vmul bf16: round-to-nearest product computed in bf16 precision."""

170 + prod = (x_bf16_fp32.astype(bfloat16) * scale_bf16_fp32.astype(bfloat16)).astype(np.float32)

171 + return prod

172 +

173 +

174 +def _fp32_to_bf16_round(x_fp32):

175 + """Hardware vcvt f32->bf16 (ROUND_R, ties-to-even) via ml_dtypes."""

176 + return x_fp32.astype(bfloat16)

177 +

178 +

179 +def fp4_pack(codes_uint8, m, n_pad):

180 + """Pack E2M1 codes 2-per-byte: low nibble = even element, high nibble = odd."""

181 + codes = codes_uint8.reshape(m, n_pad).astype(np.uint8)

182 + low = codes[:, 0::2]

183 + high = codes[:, 1::2]

184 + packed = (low & np.uint8(0xF)) | ((high & np.uint8(0xF)) << np.uint8(4))

185 + return packed.reshape(-1)

186 +

187 +

188 +def quant_bf16_to_mxfp4_dn(src_bf16_fp32, m, n_pad):

189 + """bf16 source -> MXFP4 DN. Multiply in bf16, convert bf16->fp4."""

190 + group_max = get_group_max_dn(src_bf16_fp32, group_size=32) # [hat_m, n], fp32 (bf16-exact)

191 + e8m0, scaling_fp32, _ = bf16_maxes_to_e2m1_dn(group_max)

192 +

193 + hat_m = m // 32

194 + scaled = np.empty((m, n_pad), dtype=np.float32)

195 + for rb in range(hat_m):

196 + scaled[rb * 32 : (rb + 1) * 32, :] = _bf16_mul_round(

197 + src_bf16_fp32[rb * 32 : (rb + 1) * 32, :], scaling_fp32[rb : rb + 1, :]

198 + )

199 +

200 + # Saturate NaN -> +Inf before FP4 conversion (matches SaturateBf16NaNToPosInf).

201 + with np.errstate(invalid="ignore"):

202 + scaled = np.where(np.isnan(scaled), np.float32(np.inf), scaled)

203 + # bf16 -> fp4 (round-to-nearest, ties-to-even) via ml_dtypes float4_e2m1fn.

204 + fp4_codes = scaled.astype(bfloat16).astype(float4_e2m1fn).view(np.uint8).astype(np.uint8)

205 +

206 + e8_dn = pack_e8_dn(e8m0, hat_m, n_pad, n_pad)

207 + fp4_nd = fp4_pack(fp4_codes, m, n_pad)

208 + return fp4_nd, e8_dn, group_max

209 +

210 +

211 +def fp16_maxes_to_e2m1_dn(group_max_fp16):

212 + """OCP MXFP4 e8m0 + bf16 scaling derived from the fp16 group-max BITS.

213 +

214 + The DN path (AbsReduceMax_DN) reduces fp16 abs directly (vabs on half), so the

215 + max is fp16-typed. ComputeB16OcpExponentAndScaling then reads it with the bf16

216 + exponent mask (0x7F80 >> 7) on the raw bits -- it does not reinterpret them as

217 + fp16. We mirror that bit-level extraction here.

218 + """

219 + bits = group_max_fp16.view(np.uint16).astype(np.uint16)

220 + exp_bf16 = bits & np.uint16(0x7F80)

221 + mant_bf16 = bits & np.uint16(0x007F)

222 + nan_mask = (exp_bf16 == 0x7F80) & (mant_bf16 != 0)

223 + exp_clamped = np.maximum(exp_bf16, np.uint16(0x0100))

224 + shared_exp_bits = exp_clamped.astype(np.uint16) - np.uint16(0x0100)

225 + e8m0 = ((shared_exp_bits >> 7) & np.uint16(0xFF)).astype(np.uint8)

226 + scale_bits = np.uint16(0x7F00) - shared_exp_bits

227 + scale_bits = np.where(nan_mask, np.uint16(0x7FC0), scale_bits)

228 + e8m0 = np.where(nan_mask, np.uint8(0xFF), e8m0)

229 + scaling_fp32 = (scale_bits.astype(np.uint32) << 16).view(np.float32)

230 + return e8m0, scaling_fp32

231 +

232 +

233 +def quant_fp16_to_mxfp4_dn(src_fp16, m, n_pad):

234 + """fp16 source -> MXFP4 DN. Multiply in fp32, fp32->bf16(round)->bf16->fp4.

235 +

236 + The DN reducer (AbsReduceMax_DN) takes the fp16 abs max directly (NOT a

237 + bf16-truncated max like the non-DN path), and stores it as fp16 bits. The

238 + exponent/scaling stage then reads those bits with the bf16 exponent mask.

239 + """

240 + src_fp32 = src_fp16.astype(np.float32)

241 + # DN group max: fp16 abs reduced directly, kept in fp16 precision.

242 + group_max_fp16 = np.zeros((m // 32, n_pad), dtype=np.float16)

243 + for rb in range(m // 32):

244 + group_max_fp16[rb] = np.max(np.abs(src_fp16[rb * 32 : (rb + 1) * 32, :]), axis=0)

245 + e8m0, scaling_fp32 = fp16_maxes_to_e2m1_dn(group_max_fp16)

246 +

247 + hat_m = m // 32

248 + scaled_bf16 = np.empty((m, n_pad), dtype=bfloat16)

249 + for rb in range(hat_m):

250 + # fp16 -> f32 (lossless), f32 * scale(f32), f32 -> bf16 (round-to-nearest).

251 + # Values beyond bf16 max saturate to +Inf (hardware RS_ENABLE), which the

252 + # NaN->+Inf step below then routes to the largest FP4 code.

253 + block = src_fp32[rb * 32 : (rb + 1) * 32, :]

254 + prod_fp32 = block * scaling_fp32[rb : rb + 1, :].astype(np.float32)

255 + with np.errstate(over="ignore"):

256 + scaled_bf16[rb * 32 : (rb + 1) * 32, :] = _fp32_to_bf16_round(prod_fp32)

257 +

258 + scaled_fp32 = scaled_bf16.astype(np.float32)

259 + with np.errstate(invalid="ignore"):

260 + scaled_fp32 = np.where(np.isnan(scaled_fp32), np.float32(np.inf), scaled_fp32)

261 + fp4_codes = scaled_bf16.astype(float4_e2m1fn).view(np.uint8).astype(np.uint8)

atomgit-bot6月25日

🟡 Medium Priority

在 quant_fp16_to_mxfp4_dn 函数中（文件 gen_data.py），NaN→+Inf 的处理被应用到变量 scaled_fp32 上，但随后转换为 fp4 时使用的是未经 NaN 处理的 scaled_bf16：

硬件路径 calcQuantizedFP4E2M1Values_DN_Fp16_Row 在 vcvt bf16→fp4 之前调用了 SaturateBf16NaNToPosInf。因此若输入包含 NaN，golden 与硬件输出将不一致。

触发条件：源数据中包含 NaN（当前测试数据 _gen_src_fp16_safe 生成的值在正常范围内，NaN 概率极低，故降低 confidence）。

改动建议

+scaled_fp32 = scaled_bf16.astype(np.float32)
+with np.errstate(invalid="ignore"):
+    scaled_fp32 = np.where(np.isnan(scaled_fp32), np.float32(np.inf), scaled_fp32)
-    fp4_codes = scaled_bf16.astype(float4_e2m1fn).view(np.uint8).astype(np.uint8)
+fp4_codes = scaled_fp32.astype(bfloat16).astype(float4_e2m1fn).view(np.uint8).astype(np.uint8)

应用建议

不准确？

262 +

263 + e8_dn = pack_e8_dn(e8m0, hat_m, n_pad, n_pad)

264 + fp4_nd = fp4_pack(fp4_codes, m, n_pad)

265 + return fp4_nd, e8_dn, group_max_fp16

266 +

267 +

135 CASE_PARAMS = [ 268 CASE_PARAMS = [

136 ("TQUANTDNTest.case_bf16_64x64", 64, 64), 269 ("TQUANTDNTest.case_bf16_64x64", 64, 64),

137 ("TQUANTDNTest.case_bf16_128x64", 128, 64), 270 ("TQUANTDNTest.case_bf16_128x64", 128, 64),

@@ -149,6 +282,19 @@ FP32_CASE_PARAMS = [

149 ("TQUANTDNTest.case_fp32_64x256", 64, 256), 282 ("TQUANTDNTest.case_fp32_64x256", 64, 256),

150 ] 283 ]

151 284

285 +# MXFP4 (E2M1) DN: 3 shapes each for bf16 and fp16 sources.

286 +MXFP4_BF16_CASE_PARAMS = [

287 + ("TQUANTDNTest.case_mxfp4_bf16_64x128", 64, 128),

288 + ("TQUANTDNTest.case_mxfp4_bf16_128x128", 128, 128),

289 + ("TQUANTDNTest.case_mxfp4_bf16_64x256", 64, 256),

290 +]

291 +

292 +MXFP4_FP16_CASE_PARAMS = [

293 + ("TQUANTDNTest.case_mxfp4_fp16_64x128", 64, 128),

294 + ("TQUANTDNTest.case_mxfp4_fp16_128x128", 128, 128),

295 + ("TQUANTDNTest.case_mxfp4_fp16_64x256", 64, 256),

296 +]

297 +

152 GOLDEN_DIR = os.environ.get("PTO_GOLDEN_DIR", ".") 298 GOLDEN_DIR = os.environ.get("PTO_GOLDEN_DIR", ".")

153 299

154 300

@@ -205,6 +351,65 @@ def gen_golden_data_fp32(case_name, m, n):

205 _write_golden(out_dir, input_bytes, fp8_nd, e8_dn, golden_group_max_f32.reshape(-1).tobytes()) 351 _write_golden(out_dir, input_bytes, fp8_nd, e8_dn, golden_group_max_f32.reshape(-1).tobytes())

206 352

207 353

354 +def _write_golden_mxfp4(out_dir, input_bytes, fp4_nd, e8_dn, group_max_bytes):

355 + os.makedirs(out_dir, exist_ok=True)

356 + with open(os.path.join(out_dir, "input.bin"), "wb") as f:

357 + f.write(input_bytes)

358 + with open(os.path.join(out_dir, "golden_fp4_nd.bin"), "wb") as f:

359 + f.write(fp4_nd.tobytes())

360 + with open(os.path.join(out_dir, "golden_e8_dn.bin"), "wb") as f:

361 + f.write(e8_dn.tobytes())

362 + with open(os.path.join(out_dir, "golden_group_max.bin"), "wb") as f:

363 + f.write(group_max_bytes)

364 +

365 +

366 +def gen_golden_data_mxfp4_bf16(case_name, m, n):

367 + n_pad = n

368 + src = _gen_src(m, n_pad)

369 + bf16_bits = fp32_to_bf16_bits(src).reshape(m, n_pad)

370 + src_bf16_fp32 = bf16_bits_to_fp32(bf16_bits.flatten()).reshape(m, n_pad)

371 +

372 + fp4_nd, e8_dn, group_max = quant_bf16_to_mxfp4_dn(src_bf16_fp32, m, n_pad)

373 + golden_group_max_bf16 = fp32_to_bf16_bits(group_max)

374 +

375 + out_dir = os.path.join(GOLDEN_DIR, case_name)

376 + _write_golden_mxfp4(

377 + out_dir, bf16_bits.reshape(-1).tobytes(), fp4_nd, e8_dn, golden_group_max_bf16.reshape(-1).tobytes()

378 + )

379 +

380 +

381 +def _gen_src_fp16_safe(m, n_pad):

382 + """fp16-safe variant of _gen_src: per-group max in [0.25, 16] (no *10000), so

383 + all values stay well within fp16 normal range. Uses the same RNG stream as

384 + _gen_src so bf16 tests (which use _gen_src) keep their existing goldens."""

atomgit-bot6月25日

🔵 Low Priority

_gen_src_fp16_safe 的文档字符串声称 "Uses the same RNG stream as _gen_src so bf16 tests (which use _gen_src) keep their existing goldens."，但该函数并未调用 _gen_src 或复用其 RNG 状态，而是独立调用 np.random.uniform 生成随机数。该注释具有误导性。

实际上 bf16 golden 不受影响仅仅是因为 _gen_src_fp16_safe 在 __main__ 中的调用顺序排在所有 bf16 测试之后，而非因为它使用了"相同的 RNG 流"。

这是一个文档/注释准确性问题，不影响代码正确性。

建议：修正 docstring 为更准确的描述，例如："Generates fp16-safe data with per-group max in [0.25, 16] (no *10000 multiplier). Called after all bf16 cases so existing bf16 goldens are unaffected."

不准确？

385 + hat_m = m // 32

386 + log_min = np.log2(0.25)

387 + log_max = np.log2(16.0)

388 + log_group_max = np.random.uniform(log_min, log_max, size=(hat_m, n_pad))

389 + group_max_target = (2.0**log_group_max).astype(np.float32)

390 + base = np.random.uniform(0.1, 1.0, size=(m, n_pad)).astype(np.float32)

391 + group_max_repeated = np.repeat(group_max_target, 32, axis=0)[:m, :]

392 + return base * group_max_repeated

393 +

394 +

395 +def gen_golden_data_mxfp4_fp16(case_name, m, n):

396 + n_pad = n

397 + # fp16 source: keep magnitudes in fp16 normal range (no overflow to inf), since

398 + # the DN fp16 path reduces fp16 abs directly and inf edge cases diverge.

399 + src = _gen_src_fp16_safe(m, n_pad).astype(np.float16)

400 + src_fp16_bits = src.view(np.uint16)

401 +

402 + fp4_nd, e8_dn, group_max = quant_fp16_to_mxfp4_dn(src, m, n_pad)

403 + # DN reduces fp16 abs directly and stores the max as fp16 bits (kernel MaxTile

404 + # is half when T=half), so the golden group_max is fp16 bits too.

405 + golden_group_max_fp16 = group_max.view(np.uint16)

406 +

407 + out_dir = os.path.join(GOLDEN_DIR, case_name)

408 + _write_golden_mxfp4(

409 + out_dir, src_fp16_bits.reshape(-1).tobytes(), fp4_nd, e8_dn, golden_group_max_fp16.reshape(-1).tobytes()

410 + )

411 +

412 +

208 if __name__ == "__main__": 413 if __name__ == "__main__":

209 np.random.seed(42) 414 np.random.seed(42)

210 for case_name, m, n in CASE_PARAMS: 415 for case_name, m, n in CASE_PARAMS:

@@ -213,4 +418,10 @@ if __name__ == "__main__":

213 for case_name, m, n in FP32_CASE_PARAMS: 418 for case_name, m, n in FP32_CASE_PARAMS:

214 print(f"Generating {case_name}...") 419 print(f"Generating {case_name}...")

215 gen_golden_data_fp32(case_name, m, n) 420 gen_golden_data_fp32(case_name, m, n)

421 + for case_name, m, n in MXFP4_BF16_CASE_PARAMS:

422 + print(f"Generating {case_name}...")

423 + gen_golden_data_mxfp4_bf16(case_name, m, n)

424 + for case_name, m, n in MXFP4_FP16_CASE_PARAMS:

425 + print(f"Generating {case_name}...")

426 + gen_golden_data_mxfp4_fp16(case_name, m, n)

216 print("Done.") 427 print("Done.")

Mtests/npu/a5/src/st/testcase/tquant_dn/main.cpp+200-0

	@@ -25,6 +25,12 @@ template <int Stage, int M, int N, int N_pad>
25	void LaunchTQuantDN_fp32(uint32_t src, int8_t fp8_nd, uint8_t e8_dn, int8_t fp8_nz, uint8_t *e8_zz,	25	void LaunchTQuantDN_fp32(uint32_t src, int8_t fp8_nd, uint8_t e8_dn, int8_t fp8_nz, uint8_t *e8_zz,
26	uint32_t max_dn, void stream);	26	uint32_t max_dn, void stream);
27		27
		28	+template <int M, int N, int N_pad>
		29	+void LaunchTQuantDN_MXFP4_bf16(uint16_t src, uint8_t fp4_nd, uint8_t e8_dn, uint16_t max_dn, void *stream);
		30	+
		31	+template <int M, int N, int N_pad>
		32	+void LaunchTQuantDN_MXFP4_fp16(uint16_t src, uint8_t fp4_nd, uint8_t e8_dn, uint16_t max_dn, void *stream);
		33	+
28	} // namespace TQuantDNTest	34	} // namespace TQuantDNTest
29		35
30	class TQUANTDNTest : public testing::Test {	36	class TQUANTDNTest : public testing::Test {
	@@ -298,3 +304,197 @@ TEST_F(TQUANTDNTest, case_fp32_64x256)
298	{	304	{
299	test_tquant_dn_fp32<64, 256, 256>();	305	test_tquant_dn_fp32<64, 256, 256>();
300	}	306	}
		307	+
		308	+// MXFP4 (E2M1) DN tests. Both bf16 and fp16 sources share the same UB/GM shape:
		309	+// fp4_nd : M * packedCols bytes (packedCols = paddedCols/2)
		310	+// e8_dn : hatM * paddedCols bytes
		311	+// max_dn : hatM * paddedCols * sizeof(b16) bytes
		312	+template <int M, int N, int N_pad>
		313	+void test_tquant_dn_mxfp4_bf16()
		314	+{
		315	+ constexpr int grpSize = 32;
		316	+ constexpr int hatM = M / grpSize;
		317	+ constexpr int paddedCols = N_pad;
		318	+ constexpr int packedCols = paddedCols / 2;
		319	+ size_t srcFileSize = M * paddedCols * sizeof(uint16_t);
		320	+ size_t fp4NDFileSize = M * packedCols * sizeof(uint8_t);
		321	+ size_t e8DNFileSize = hatM * paddedCols * sizeof(uint8_t);
		322	+ size_t maxDNFileSize = hatM * paddedCols * sizeof(uint16_t);
		323	+
		324	+ aclInit(nullptr);
		325	+ aclrtSetDevice(0);
		326	+ aclrtStream stream;
		327	+ aclrtCreateStream(&stream);
		328	+
		329	+ uint8_t *srcHost;
		330	+ uint8_t *fp4NDHost;
		331	+ uint8_t *e8DNHost;
		332	+ uint16_t *maxDNHost;
		333	+ uint8_t *srcDevice;
		334	+ uint8_t *fp4NDDevice;
		335	+ uint8_t *e8DNDevice;
		336	+ uint16_t *maxDNDevice;
		337	+
		338	+ aclrtMallocHost((void **)(&srcHost), srcFileSize);
		339	+ aclrtMallocHost((void **)(&fp4NDHost), fp4NDFileSize);
		340	+ aclrtMallocHost((void **)(&e8DNHost), e8DNFileSize);
		341	+ aclrtMallocHost((void **)(&maxDNHost), maxDNFileSize);
		342	+
		343	+ aclrtMalloc((void **)&srcDevice, srcFileSize, ACL_MEM_MALLOC_HUGE_FIRST);
		344	+ aclrtMalloc((void **)&fp4NDDevice, fp4NDFileSize, ACL_MEM_MALLOC_HUGE_FIRST);
		345	+ aclrtMalloc((void **)&e8DNDevice, e8DNFileSize, ACL_MEM_MALLOC_HUGE_FIRST);
		346	+ aclrtMalloc((void **)&maxDNDevice, maxDNFileSize, ACL_MEM_MALLOC_HUGE_FIRST);
		347	+
		348	+ ReadFile(GetGoldenDir() + "/input.bin", srcFileSize, srcHost, srcFileSize);
		349	+ aclrtMemcpy(srcDevice, srcFileSize, srcHost, srcFileSize, ACL_MEMCPY_HOST_TO_DEVICE);
		350	+
		351	+ const std::string goldenDir = GetGoldenDir();
		352	+
		353	+ TQuantDNTest::LaunchTQuantDN_MXFP4_bf16<M, N, N_pad>((uint16_t *)srcDevice, fp4NDDevice, e8DNDevice, maxDNDevice,
		354	+ stream);
		355	+ aclError syncRet = aclrtSynchronizeStream(stream);
		356	+ ASSERT_EQ(syncRet, ACL_SUCCESS) << "MXFP4 bf16 DN sync failed: " << aclGetRecentErrMsg();
		357	+
		358	+ aclrtMemcpy(fp4NDHost, fp4NDFileSize, fp4NDDevice, fp4NDFileSize, ACL_MEMCPY_DEVICE_TO_HOST);
		359	+ aclrtMemcpy(e8DNHost, e8DNFileSize, e8DNDevice, e8DNFileSize, ACL_MEMCPY_DEVICE_TO_HOST);
		360	+ aclrtMemcpy(maxDNHost, maxDNFileSize, maxDNDevice, maxDNFileSize, ACL_MEMCPY_DEVICE_TO_HOST);
		361	+ WriteFile(goldenDir + "/output_fp4_nd.bin", fp4NDHost, fp4NDFileSize);
		362	+ WriteFile(goldenDir + "/output_e8_dn.bin", e8DNHost, e8DNFileSize);
		363	+ WriteFile(goldenDir + "/output_group_max.bin", maxDNHost, maxDNFileSize);
		364	+
		365	+ std::vector<uint8_t> goldenFp4Nd(fp4NDFileSize);
		366	+ std::vector<uint8_t> goldenE8Dn(e8DNFileSize);
		367	+ std::vector<uint16_t> goldenGroupMax(maxDNFileSize / sizeof(uint16_t));
		368	+ std::vector<uint8_t> outFp4Nd(fp4NDFileSize);
		369	+ std::vector<uint8_t> outE8Dn(e8DNFileSize);
		370	+ std::vector<uint16_t> outGroupMax(maxDNFileSize / sizeof(uint16_t));
		371	+ ReadFile(goldenDir + "/golden_fp4_nd.bin", fp4NDFileSize, goldenFp4Nd.data(), fp4NDFileSize);
		372	+ ReadFile(goldenDir + "/golden_e8_dn.bin", e8DNFileSize, goldenE8Dn.data(), e8DNFileSize);
		373	+ ReadFile(goldenDir + "/golden_group_max.bin", maxDNFileSize, goldenGroupMax.data(), maxDNFileSize);
		374	+ ReadFile(goldenDir + "/output_fp4_nd.bin", fp4NDFileSize, outFp4Nd.data(), fp4NDFileSize);
		375	+ ReadFile(goldenDir + "/output_e8_dn.bin", e8DNFileSize, outE8Dn.data(), e8DNFileSize);
		376	+ ReadFile(goldenDir + "/output_group_max.bin", maxDNFileSize, outGroupMax.data(), maxDNFileSize);
		377	+ ExpectGoldenMatch("MXFP4_BF16_DN", "fp4_nd", goldenFp4Nd, outFp4Nd);
		378	+ ExpectGoldenMatch("MXFP4_BF16_DN", "e8_dn (exponents)", goldenE8Dn, outE8Dn);
		379	+ ExpectGoldenMatch("MXFP4_BF16_DN", "group_max", goldenGroupMax, outGroupMax);
		380	+
		381	+ aclrtFree(srcDevice);
		382	+ aclrtFree(fp4NDDevice);
		383	+ aclrtFree(e8DNDevice);
		384	+ aclrtFree(maxDNDevice);
		385	+ aclrtFreeHost(srcHost);
		386	+ aclrtFreeHost(fp4NDHost);
		387	+ aclrtFreeHost(e8DNHost);
		388	+ aclrtFreeHost(maxDNHost);
		389	+ aclrtDestroyStream(stream);
		390	+ aclrtResetDevice(0);
		391	+ aclFinalize();
		392	+}
		393	+
		394	+TEST_F(TQUANTDNTest, case_mxfp4_bf16_64x128)
		395	+{
		396	+ test_tquant_dn_mxfp4_bf16<64, 128, 128>();
		397	+}
		398	+TEST_F(TQUANTDNTest, case_mxfp4_bf16_128x128)
		399	+{
		400	+ test_tquant_dn_mxfp4_bf16<128, 128, 128>();
		401	+}
		402	+TEST_F(TQUANTDNTest, case_mxfp4_bf16_64x256)
		403	+{
		404	+ test_tquant_dn_mxfp4_bf16<64, 256, 256>();
		405	+}
		406	+
		407	+template <int M, int N, int N_pad>
		408	+void test_tquant_dn_mxfp4_fp16()
		409	+{
		410	+ constexpr int grpSize = 32;
		411	+ constexpr int hatM = M / grpSize;
		412	+ constexpr int paddedCols = N_pad;
		413	+ constexpr int packedCols = paddedCols / 2;
		414	+ size_t srcFileSize = M * paddedCols * sizeof(uint16_t);
		415	+ size_t fp4NDFileSize = M * packedCols * sizeof(uint8_t);
		416	+ size_t e8DNFileSize = hatM * paddedCols * sizeof(uint8_t);
		417	+ size_t maxDNFileSize = hatM * paddedCols * sizeof(uint16_t);
		418	+
		419	+ aclInit(nullptr);
		420	+ aclrtSetDevice(0);
		421	+ aclrtStream stream;
		422	+ aclrtCreateStream(&stream);
		423	+
		424	+ uint8_t *srcHost;
		425	+ uint8_t *fp4NDHost;
		426	+ uint8_t *e8DNHost;
		427	+ uint16_t *maxDNHost;
		428	+ uint8_t *srcDevice;
		429	+ uint8_t *fp4NDDevice;
		430	+ uint8_t *e8DNDevice;
		431	+ uint16_t *maxDNDevice;
		432	+
		433	+ aclrtMallocHost((void **)(&srcHost), srcFileSize);
		434	+ aclrtMallocHost((void **)(&fp4NDHost), fp4NDFileSize);
		435	+ aclrtMallocHost((void **)(&e8DNHost), e8DNFileSize);
		436	+ aclrtMallocHost((void **)(&maxDNHost), maxDNFileSize);
		437	+
		438	+ aclrtMalloc((void **)&srcDevice, srcFileSize, ACL_MEM_MALLOC_HUGE_FIRST);
		439	+ aclrtMalloc((void **)&fp4NDDevice, fp4NDFileSize, ACL_MEM_MALLOC_HUGE_FIRST);
		440	+ aclrtMalloc((void **)&e8DNDevice, e8DNFileSize, ACL_MEM_MALLOC_HUGE_FIRST);
		441	+ aclrtMalloc((void **)&maxDNDevice, maxDNFileSize, ACL_MEM_MALLOC_HUGE_FIRST);
		442	+
		443	+ ReadFile(GetGoldenDir() + "/input.bin", srcFileSize, srcHost, srcFileSize);
		444	+ aclrtMemcpy(srcDevice, srcFileSize, srcHost, srcFileSize, ACL_MEMCPY_HOST_TO_DEVICE);
		445	+
		446	+ const std::string goldenDir = GetGoldenDir();
		447	+
		448	+ TQuantDNTest::LaunchTQuantDN_MXFP4_fp16<M, N, N_pad>((uint16_t *)srcDevice, fp4NDDevice, e8DNDevice, maxDNDevice,
		449	+ stream);
		450	+ aclError syncRet = aclrtSynchronizeStream(stream);
		451	+ ASSERT_EQ(syncRet, ACL_SUCCESS) << "MXFP4 fp16 DN sync failed: " << aclGetRecentErrMsg();
		452	+
		453	+ aclrtMemcpy(fp4NDHost, fp4NDFileSize, fp4NDDevice, fp4NDFileSize, ACL_MEMCPY_DEVICE_TO_HOST);
		454	+ aclrtMemcpy(e8DNHost, e8DNFileSize, e8DNDevice, e8DNFileSize, ACL_MEMCPY_DEVICE_TO_HOST);
		455	+ aclrtMemcpy(maxDNHost, maxDNFileSize, maxDNDevice, maxDNFileSize, ACL_MEMCPY_DEVICE_TO_HOST);
		456	+ WriteFile(goldenDir + "/output_fp4_nd.bin", fp4NDHost, fp4NDFileSize);
		457	+ WriteFile(goldenDir + "/output_e8_dn.bin", e8DNHost, e8DNFileSize);
		458	+ WriteFile(goldenDir + "/output_group_max.bin", maxDNHost, maxDNFileSize);
		459	+
		460	+ std::vector<uint8_t> goldenFp4Nd(fp4NDFileSize);
		461	+ std::vector<uint8_t> goldenE8Dn(e8DNFileSize);
		462	+ std::vector<uint16_t> goldenGroupMax(maxDNFileSize / sizeof(uint16_t));
		463	+ std::vector<uint8_t> outFp4Nd(fp4NDFileSize);
		464	+ std::vector<uint8_t> outE8Dn(e8DNFileSize);
		465	+ std::vector<uint16_t> outGroupMax(maxDNFileSize / sizeof(uint16_t));
		466	+ ReadFile(goldenDir + "/golden_fp4_nd.bin", fp4NDFileSize, goldenFp4Nd.data(), fp4NDFileSize);
		467	+ ReadFile(goldenDir + "/golden_e8_dn.bin", e8DNFileSize, goldenE8Dn.data(), e8DNFileSize);
		468	+ ReadFile(goldenDir + "/golden_group_max.bin", maxDNFileSize, goldenGroupMax.data(), maxDNFileSize);
		469	+ ReadFile(goldenDir + "/output_fp4_nd.bin", fp4NDFileSize, outFp4Nd.data(), fp4NDFileSize);
		470	+ ReadFile(goldenDir + "/output_e8_dn.bin", e8DNFileSize, outE8Dn.data(), e8DNFileSize);
		471	+ ReadFile(goldenDir + "/output_group_max.bin", maxDNFileSize, outGroupMax.data(), maxDNFileSize);
		472	+ ExpectGoldenMatch("MXFP4_FP16_DN", "fp4_nd", goldenFp4Nd, outFp4Nd);
		473	+ ExpectGoldenMatch("MXFP4_FP16_DN", "e8_dn (exponents)", goldenE8Dn, outE8Dn);
		474	+ ExpectGoldenMatch("MXFP4_FP16_DN", "group_max", goldenGroupMax, outGroupMax);
		475	+
		476	+ aclrtFree(srcDevice);
		477	+ aclrtFree(fp4NDDevice);
		478	+ aclrtFree(e8DNDevice);
		479	+ aclrtFree(maxDNDevice);
		480	+ aclrtFreeHost(srcHost);
		481	+ aclrtFreeHost(fp4NDHost);
		482	+ aclrtFreeHost(e8DNHost);
		483	+ aclrtFreeHost(maxDNHost);
		484	+ aclrtDestroyStream(stream);
		485	+ aclrtResetDevice(0);
		486	+ aclFinalize();
		487	+}
		488	+
		489	+TEST_F(TQUANTDNTest, case_mxfp4_fp16_64x128)
		490	+{
		491	+ test_tquant_dn_mxfp4_fp16<64, 128, 128>();
		492	+}
		493	+TEST_F(TQUANTDNTest, case_mxfp4_fp16_128x128)
		494	+{
		495	+ test_tquant_dn_mxfp4_fp16<128, 128, 128>();
		496	+}
		497	+TEST_F(TQUANTDNTest, case_mxfp4_fp16_64x256)
		498	+{
		499	+ test_tquant_dn_mxfp4_fp16<64, 256, 256>();
		500	+}

Mtests/npu/a5/src/st/testcase/tquant_dn/tquant_dn_kernel.cpp+150-1

	@@ -19,6 +19,32 @@ using namespace pto;
19		19
20	namespace TQuantDNTest {	20	namespace TQuantDNTest {
21		21
		22	+// Copy packed FP4 rows from the TQUANT output UB region (2D [rows, srcStride]) into
		23	+// the TSTORE tile UB region (2D [rows, dstStride]), honouring validPackedCols. Mirrors
		24	+// the proven CompactFp4PackedRows helper used by the non-DN MXFP4 kernel.
		25	+PTO_INTERNAL void CompactFp4PackedRows(__ubuf__ uint8_t dstPtr, __ubuf__ uint8_t srcPtr, uint32_t rows,
		26	+ uint32_t validPackedCols, uint32_t srcStride, uint32_t dstStride)
		27	+{
		28	+ constexpr uint32_t elementsPerRepeat = REPEAT_BYTE / sizeof(uint8_t);
		29	+ RegTensor<uint8_t> vreg;
		30	+ UnalignReg ureg;
		31	+ uint16_t repeatTimes = CeilDivision(validPackedCols, elementsPerRepeat);
		32	+ for (uint16_t row = 0; row < (uint16_t)rows; ++row) {
		33	+ uint32_t remaining = validPackedCols;
		34	+ __ubuf__ uint8_t srcRow = srcPtr + row srcStride;
		35	+ __ubuf__ uint8_t dstRow = dstPtr + row dstStride;
		36	+ for (uint16_t repeat = 0; repeat < repeatTimes; ++repeat) {
		37	+ uint32_t cols = remaining > elementsPerRepeat ? elementsPerRepeat : remaining;
		38	+ MaskReg preg = CreatePredicate<uint8_t>(cols);
		39	+ uint32_t offset = repeat * elementsPerRepeat;
		40	+ vldas(ureg, srcRow + offset);
		41	+ vldus(vreg, ureg, srcRow + offset);
		42	+ vsts(vreg, dstRow, offset, NORM_B8, preg);
		43	+ remaining -= cols;
		44	+ }
		45	+ }
		46	+}
		47	+
22	// Stage 1: after TQUANT (FP8 ND + E8M0 DN)	48	// Stage 1: after TQUANT (FP8 ND + E8M0 DN)
23	// Stage 2: after TQUANT + FP8 ND->NZ	49	// Stage 2: after TQUANT + FP8 ND->NZ
24	// Stage 3: full pipeline including E8 DN->ZZ	50	// Stage 3: full pipeline including E8 DN->ZZ
	@@ -151,7 +177,9 @@ __global__ AICORE void runTQuantDN(__gm__ T __in__ *src_gm, __gm__ int8_t __out_
151	set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);	177	set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);
152	wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);	178	wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);
153		179
154	- TQUANT<QuantType::MXFP8, QuantScaleAlg::OCP>(fp8Tile, srcTile, &e8Tile, &maxPerGpTile, &scalingTile, &e8DnTile);	180	+ // Generic DN API: grp_axis=0 (groups on axis 0), single MxQuantAlg tag. The
		181	+ // exponent is written into e8Tile; the caller reshapes it into e8DnTile via TMOV.
		182	+ TQuant<0, MxQuantAlg::OcpMxFp8E4M3>(fp8Tile, srcTile, &e8Tile, &maxPerGpTile, &scalingTile);
155		183
156	// TQuant writes the exponent tile (e8Tile) in row-major [hatM, paddedCols].	184	// TQuant writes the exponent tile (e8Tile) in row-major [hatM, paddedCols].
157	// Copy it to e8DnTile so the UB tile shape matches the GM shape exactly.	185	// Copy it to e8DnTile so the UB tile shape matches the GM shape exactly.
	@@ -199,6 +227,114 @@ void LaunchTQuantDN_fp32(uint32_t src, int8_t fp8_nd, uint8_t e8_dn, int8_t
199	<<<1, nullptr, stream>>>((float )src, fp8_nd, e8_dn, fp8_nz, e8_zz, (float )max_dn);	227	<<<1, nullptr, stream>>>((float )src, fp8_nd, e8_dn, fp8_nz, e8_zz, (float )max_dn);
200	}	228	}
201		229
		230	+// MXFP4 (E2M1) DN kernel: quantizes src[M,N_pad] to packed FP4 plus per-group
		231	+// e8m0/max tiles. Stage 3 writes FP4 into UB at byte offset (r * StaticCols + off)/2,
		232	+// so the FP4 region is a 2D [M, packedCols] layout (packedCols = N_pad/2). A flat
		233	+// float4_e2m1x2_t tile is the TQUANT output; CompactFp4PackedRows bridges it to a
		234	+// uint8_t tile used by TSTORE (proven non-DN MXFP4 store pattern).
		235	+template <typename T, int M, int N, int N_pad>
		236	+__global__ AICORE void runTQuantDN_MXFP4(__gm__ T __in__ src_gm, __gm__ uint8_t __out__ fp4_nd_gm,
		237	+ __gm__ uint8_t __out__ e8_dn_gm, __gm__ T __out__ max_dn_gm)
		238	+{
		239	+ constexpr uint32_t grpSize = 32;
		240	+ constexpr uint32_t hatM = M / grpSize;
		241	+ constexpr uint32_t paddedCols = N_pad;
		242	+ constexpr uint32_t packedCols = paddedCols / 2;
		243	+ constexpr uint32_t fp4FlatAligned = PTO_CEIL(M * packedCols, 32);
		244	+
		245	+ using SrcTile =
		246	+ Tile<TileType::Vec, T, M, paddedCols, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Zero>;
		247	+ using MaxTile =
		248	+ Tile<TileType::Vec, T, hatM, paddedCols, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Zero>;
		249	+ using ScalingTile =
		250	+ Tile<TileType::Vec, T, hatM, paddedCols, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Zero>;
		251	+ using E8Tile = Tile<TileType::Vec, uint8_t, hatM, paddedCols, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512,
		252	+ PadValue::Zero>;
		253	+ // TQUANT output tile: flat float4_e2m1x2_t (element = 0.5 byte -> bytes = M*packedCols).
		254	+ using DstFP4Tile = Tile<TileType::Vec, float4_e2m1x2_t, 1, fp4FlatAligned, BLayout::RowMajor, -1, -1,
		255	+ SLayout::NoneBox, 512, PadValue::Zero>;
		256	+ // TSTORE tile: uint8_t [M, packedCols], same UB region bridged by CompactFp4PackedRows.
		257	+ // Valid extents are DYNAMIC so the tile can be runtime-sized to (M, packedCols).
		258	+ using DstBytesTile =
		259	+ Tile<TileType::Vec, uint8_t, M, packedCols, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Zero>;
		260	+
		261	+ constexpr uint32_t srcBytes = M * paddedCols * sizeof(T);
		262	+ constexpr uint32_t maxBytes = hatM * paddedCols * sizeof(T);
		263	+ constexpr uint32_t scalingBytes = hatM * paddedCols * sizeof(T);
		264	+ constexpr uint32_t e8Bytes = hatM * paddedCols;
		265	+ constexpr uint32_t fp4Bytes = M * packedCols;
		266	+
		267	+ constexpr uint32_t srcAddr = 0x0;
		268	+ constexpr uint32_t maxAddr = PTO_CEIL(srcAddr + srcBytes, 0x20);
		269	+ constexpr uint32_t scalingAddr = PTO_CEIL(maxAddr + maxBytes, 0x20);
		270	+ constexpr uint32_t e8Addr = PTO_CEIL(scalingAddr + scalingBytes, 0x20);
		271	+ constexpr uint32_t fp4Addr = PTO_CEIL(e8Addr + e8Bytes, 0x20);
		272	+ constexpr uint32_t fp4StoreAddr = PTO_CEIL(fp4Addr + fp4Bytes, 0x20);
		273	+ constexpr uint32_t layoutEnd = PTO_CEIL(fp4StoreAddr + fp4Bytes, 0x100);
		274	+ static_assert(layoutEnd <= 0x40000, "MXFP4 DN UB layout exceeds 256 KB.");
		275	+
		276	+ SrcTile srcTile(M, paddedCols);
		277	+ MaxTile maxTile(hatM, paddedCols);
		278	+ ScalingTile scalingTile(hatM, paddedCols);
		279	+ E8Tile e8Tile(hatM, paddedCols);
		280	+ DstFP4Tile fp4Tile;
		281	+ DstBytesTile fp4BytesTile(M, packedCols);
		282	+
		283	+ TASSIGN(srcTile, srcAddr);
		284	+ TASSIGN(maxTile, maxAddr);
		285	+ TASSIGN(scalingTile, scalingAddr);
		286	+ TASSIGN(e8Tile, e8Addr);
		287	+ TASSIGN(fp4Tile, fp4Addr);
		288	+ TASSIGN(fp4BytesTile, fp4StoreAddr);
		289	+
		290	+ using SrcGlobal = GlobalTensor<T, Shape<1, 1, 1, M, N_pad>, pto::Stride<1, 1, 1, N_pad, 1>>;
		291	+ using Fp4Global = GlobalTensor<uint8_t, Shape<1, 1, 1, M, packedCols>, pto::Stride<1, 1, 1, packedCols, 1>>;
		292	+ using MaxGlobal = GlobalTensor<T, Shape<1, 1, 1, hatM, paddedCols>, pto::Stride<1, 1, 1, paddedCols, 1>>;
		293	+ using E8Global = GlobalTensor<uint8_t, Shape<1, 1, 1, hatM, paddedCols>, pto::Stride<1, 1, 1, paddedCols, 1>>;
		294	+
		295	+ SrcGlobal srcGlobal(src_gm);
		296	+ Fp4Global fp4Global(fp4_nd_gm);
		297	+ MaxGlobal maxGlobal(max_dn_gm);
		298	+ E8Global e8Global(e8_dn_gm);
		299	+
		300	+ TLOAD(srcTile, srcGlobal);
		301	+ set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);
		302	+ wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);
		303	+
		304	+ // Generic DN API: grp_axis=0, single MxQuantAlg tag. TQuant writes the exponent
		305	+ // into e8Tile; the kernel TSTOREs e8Tile directly (shape already matches GM).
		306	+ TQuant<0, MxQuantAlg::OcpMxFp4E2M1>(fp4Tile, srcTile, &e8Tile, &maxTile, &scalingTile);
		307	+
		308	+ __VEC_SCOPE__
		309	+ {
		310	+ // Stage 3 wrote FP4 densely at [r, packedCols] (srcStride == dstStride == packedCols),
		311	+ // so the compact copy is a plain row-wise move into the TSTORE tile's UB region.
		312	+ mem_bar(VST_VLD);
		313	+ CompactFp4PackedRows((__ubuf__ uint8_t )fp4BytesTile.data(), (__ubuf__ uint8_t )fp4Tile.data(), M, packedCols,
		314	+ packedCols, packedCols);
		315	+ mem_bar(VST_VST);
		316	+ }
		317	+
		318	+ set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0);
		319	+ wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0);
		320	+ TSTORE(fp4Global, fp4BytesTile);
		321	+ TSTORE(maxGlobal, maxTile);
		322	+ TSTORE(e8Global, e8Tile);
		323	+}
		324	+
		325	+template <int M, int N, int N_pad>
		326	+void LaunchTQuantDN_MXFP4_bf16(uint16_t src, uint8_t fp4_nd, uint8_t e8_dn, uint16_t max_dn, void *stream)
		327	+{
		328	+ runTQuantDN_MXFP4<bfloat16_t, M, N, N_pad>
		329	+ <<<1, nullptr, stream>>>((bfloat16_t )src, fp4_nd, e8_dn, (bfloat16_t )max_dn);
		330	+}
		331	+
		332	+template <int M, int N, int N_pad>
		333	+void LaunchTQuantDN_MXFP4_fp16(uint16_t src, uint8_t fp4_nd, uint8_t e8_dn, uint16_t max_dn, void *stream)
		334	+{
		335	+ runTQuantDN_MXFP4<half, M, N, N_pad><<<1, nullptr, stream>>>((half )src, fp4_nd, e8_dn, (half )max_dn);
		336	+}
		337	+
202	#define INSTANTIATE_TQUANT_DN_STAGE(S, M, N, NP) \	338	#define INSTANTIATE_TQUANT_DN_STAGE(S, M, N, NP) \
203	template void LaunchTQuantDN<S, M, N, NP>(uint16_t , int8_t , uint8_t , int8_t , uint8_t , uint16_t , void *)	339	template void LaunchTQuantDN<S, M, N, NP>(uint16_t , int8_t , uint8_t , int8_t , uint8_t , uint16_t , void *)
204		340
	@@ -218,6 +354,19 @@ INSTANTIATE_TQUANT_DN_STAGE_FP32(1, 64, 128, 128);
218	INSTANTIATE_TQUANT_DN_STAGE_FP32(1, 128, 128, 128);	354	INSTANTIATE_TQUANT_DN_STAGE_FP32(1, 128, 128, 128);
219	INSTANTIATE_TQUANT_DN_STAGE_FP32(1, 64, 256, 256);	355	INSTANTIATE_TQUANT_DN_STAGE_FP32(1, 64, 256, 256);
220		356
		357	+#define INSTANTIATE_TQUANT_DN_MXFP4_BF16(M, N, NP) \
		358	+ template void LaunchTQuantDN_MXFP4_bf16<M, N, NP>(uint16_t , uint8_t , uint8_t , uint16_t , void *)
		359	+
		360	+#define INSTANTIATE_TQUANT_DN_MXFP4_FP16(M, N, NP) \
		361	+ template void LaunchTQuantDN_MXFP4_fp16<M, N, NP>(uint16_t , uint8_t , uint8_t , uint16_t , void *)
		362	+
		363	+INSTANTIATE_TQUANT_DN_MXFP4_BF16(64, 128, 128);
		364	+INSTANTIATE_TQUANT_DN_MXFP4_BF16(128, 128, 128);
		365	+INSTANTIATE_TQUANT_DN_MXFP4_BF16(64, 256, 256);
		366	+INSTANTIATE_TQUANT_DN_MXFP4_FP16(64, 128, 128);
		367	+INSTANTIATE_TQUANT_DN_MXFP4_FP16(128, 128, 128);
		368	+INSTANTIATE_TQUANT_DN_MXFP4_FP16(64, 256, 256);
		369	+
221	#undef INSTANTIATE_TQUANT_DN_STAGE	370	#undef INSTANTIATE_TQUANT_DN_STAGE
222		371
223	} // namespace TQuantDNTest	372	} // namespace TQuantDNTest