[KMCompiler] [ttx/npu] add DeepSeek-V4 operator rms_norm_dynamic_quant

[YangLong114514]

Description

Add TLE Triton implementation of MojoRMSNormDynamicQuant for DeepSeek-V4 on Ascend 910B.

Changes

• Preloads gamma and smooth_scale into UB.
• Pre-computes fused weight gamma * smooth_scale.
• Uses DSA copy for input rows.
• Computes RMSNorm, dynamic scale, and int8 quantization in a single pass per row.
• Uses inv_scale = 127 / max_abs to avoid division in the quantization path.
• Introduced row-batched prefill kernels with manually unrolled BLOCK_M=2 and BLOCK_M=4 paths.
• Tuned BLOCK_M by hidden size to balance loop amortization and UB/register pressure.
• Added prefill grid selection to improve load balancing and avoid inefficient remainder paths.

Performance

Test environment: Ascend 910B, TLE Triton backend.(cann 9.0)

shape (BS, N, D)	ascendc (us)	non-tle (us)	tle (us)	speedup (non-tle)	speedup (tle)
(64, 1, 512)	7.44	5.40	5.22	1.38x	1.43x
(64, 1, 1024)	7.77	5.84	5.92	1.33x	1.31x
(64, 1, 1536)	8.08	5.90	5.89	1.37x	1.37x
(64, 1, 2048)	8.31	6.14	6.15	1.35x	1.35x
(64, 1, 4096)	9.55	7.77	7.80	1.23x	1.22x
(64, 1024, 512)	777.56	1096.16	1085.27	0.71x	0.72x
(64, 1024, 1024)	1039.85	1316.67	1309.43	0.79x	0.79x
(64, 1024, 1536)	1334.57	1751.06	1748.04	0.76x	0.76x
(64, 1024, 2048)	1610.98	1907.25	1900.73	0.85x	0.85x
(64, 1024, 4096)	2850.37	3266.12	3261.56	0.87x	0.87x

The above perf data was obtained from testing using the perf-rms_norm_dynamic_quant.py.py script.

Accuracy Test

All accuracy test cases passed.

pytest test_normalization.py::test_rms_norm_dynamic_quant -s

[gemini-code-assist]

Code Review

This pull request introduces MojoRMSNormDynamicQuant, a fused RMSNorm and dynamic per-token quantization operator, along with its TTX backend implementation and corresponding Triton kernels for NPU. The feedback highlights critical issues in the Triton kernels, specifically potential out-of-bounds memory writes in unmasked kernels when the hidden dimension is not aligned to BLOCK_N_ALIGN. Additionally, the reviewer recommends adding division-by-zero guards in the prefill kernels and utilizing propagate_nan=tl.PropagateNan.ALL in tl.max operations to prevent performance overhead and align with NPU hardware behavior.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.