[wip]feat(ascend): add RMSNorm Ascend C kernel - #324
Closed
erfgss wants to merge 7 commits into
Closed
Conversation
[fea][Multiplatform]: add multi-architecture build support
erfgss
requested review from
EthanZero2Hero,
Flink-ddd,
KJLdefeated,
bitborne and
inaniloquentee
as code owners
August 20, 2026 12:27
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Collaborator
|
Thanks, I will review this tmr |
zhangj1an
requested changes
Aug 21, 2026
Collaborator
There was a problem hiding this comment.
Hi @erfgss , the RMSNorm kernel can pass gtest after fixing two small parts as shown below. After you fix them the PR can be merged.
Thank you!
Below are all the tests passing after fixing these 2 parts.
Here is the result for running pytest on tests/test_rms_norm_ascend.py
Details
============================= test session starts ==============================
platform linux -- Python 3.10.14, pytest-9.1.1, pluggy-1.6.0 -- /home/ma-user/work/z84450661/venv310/bin/python
cachedir: .pytest_cache
rootdir: /home/ma-user/work/z84450661/rl-kernel
configfile: pyproject.toml
plugins: anyio-4.14.1
collecting ... collected 19 items
tests/test_rms_norm_ascend.py::test_ascend_backward_formula_matches_autograd_on_cpu PASSED [ 5%]
tests/test_rms_norm_ascend.py::test_ascend_candidate_is_registered_for_rmsnorm_and_qk_norm PASSED [ 10%]
tests/test_rms_norm_ascend.py::test_ascend_forward_matches_reference[shape0-dtype0] PASSED [ 15%]
tests/test_rms_norm_ascend.py::test_ascend_forward_matches_reference[shape0-dtype1] PASSED [ 21%]
tests/test_rms_norm_ascend.py::test_ascend_forward_matches_reference[shape0-dtype2] PASSED [ 26%]
tests/test_rms_norm_ascend.py::test_ascend_forward_matches_reference[shape1-dtype0] PASSED [ 31%]
tests/test_rms_norm_ascend.py::test_ascend_forward_matches_reference[shape1-dtype1] PASSED [ 36%]
tests/test_rms_norm_ascend.py::test_ascend_forward_matches_reference[shape1-dtype2] PASSED [ 42%]
tests/test_rms_norm_ascend.py::test_ascend_forward_matches_reference[shape2-dtype0] PASSED [ 47%]
tests/test_rms_norm_ascend.py::test_ascend_forward_matches_reference[shape2-dtype1] PASSED [ 52%]
tests/test_rms_norm_ascend.py::test_ascend_forward_matches_reference[shape2-dtype2] PASSED [ 57%]
tests/test_rms_norm_ascend.py::test_ascend_forward_matches_reference[shape3-dtype0] PASSED [ 63%]
tests/test_rms_norm_ascend.py::test_ascend_forward_matches_reference[shape3-dtype1] PASSED [ 68%]
tests/test_rms_norm_ascend.py::test_ascend_forward_matches_reference[shape3-dtype2] PASSED [ 73%]
tests/test_rms_norm_ascend.py::test_ascend_backward_matches_reference[dtype0] PASSED [ 78%]
tests/test_rms_norm_ascend.py::test_ascend_backward_matches_reference[dtype1] PASSED [ 84%]
tests/test_rms_norm_ascend.py::test_ascend_forward_is_batch_and_position_invariant PASSED [ 89%]
tests/test_rms_norm_ascend.py::test_ascend_forward_fp32_and_empty_batch PASSED [ 94%]
tests/test_rms_norm_ascend.py::test_ascend_large_hidden_falls_back_to_native PASSED [100%]
=============================== warnings summary ===============================
../venv310/lib/python3.10/site-packages/torch/jit/_script.py:362: 14 warnings
/home/ma-user/work/z84450661/venv310/lib/python3.10/site-packages/torch/jit/_script.py:362: DeprecationWarning: `torch.jit.script_method` is deprecated. Please switch to `torch.compile` or `torch.export`.
warnings.warn(
tests/test_rms_norm_ascend.py::test_ascend_forward_matches_reference[shape0-dtype0]
/home/ma-user/work/z84450661/venv310/lib/python3.10/site-packages/torch/autograd/function.py:583: UserWarning: Cannot create tensor with interal format while allow_internel_format=False, tensor will be created with base format. (Triggered internally at ../torch_npu/csrc/aten/common/TensorFactories.cpp:340.)
return super().apply(*args, **kwargs) # type: ignore[misc]
-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
======================= 19 passed, 15 warnings in 5.04s ========================
Here is the result for running pytest on tests/test_rms_norm.py and rl_engine/tests/test_dispatch.py
Details
============================= test session starts ==============================
platform linux -- Python 3.10.14, pytest-9.1.1, pluggy-1.6.0 -- /home/ma-user/work/z84450661/venv310/bin/python
cachedir: .pytest_cache
rootdir: /home/ma-user/work/z84450661/rl-kernel
configfile: pyproject.toml
plugins: anyio-4.14.1
collecting ... collected 60 items
tests/test_rms_norm.py::test_forward_fp32_matches_torch_reference[4096] PASSED [ 1%]
tests/test_rms_norm.py::test_forward_fp32_matches_torch_reference[128] PASSED [ 3%]
tests/test_rms_norm.py::test_forward_fp32_matches_manual_reference[4096] PASSED [ 5%]
tests/test_rms_norm.py::test_forward_fp32_matches_manual_reference[128] PASSED [ 6%]
tests/test_rms_norm.py::test_batch_invariance_slice[4096] PASSED [ 8%]
tests/test_rms_norm.py::test_batch_invariance_slice[128] PASSED [ 10%]
tests/test_rms_norm.py::test_batch_invariance_with_padding PASSED [ 11%]
tests/test_rms_norm.py::test_dtype_paths[dtype0] PASSED [ 13%]
tests/test_rms_norm.py::test_dtype_paths[dtype1] PASSED [ 15%]
tests/test_rms_norm.py::test_dtype_paths[dtype2] PASSED [ 16%]
tests/test_rms_norm.py::test_low_precision_within_tolerance[dtype0-0.02-0.016] PASSED [ 18%]
tests/test_rms_norm.py::test_low_precision_within_tolerance[dtype1-0.001-0.001] PASSED [ 20%]
tests/test_rms_norm.py::test_eps_inside_sqrt PASSED [ 21%]
tests/test_rms_norm.py::test_weight_scaling_no_plus_one PASSED [ 23%]
tests/test_rms_norm.py::test_bad_weight_shape_raises PASSED [ 25%]
tests/test_rms_norm.py::test_inputs_not_mutated PASSED [ 26%]
tests/test_rms_norm.py::test_gradient_flows PASSED [ 28%]
tests/test_rms_norm.py::test_backward_batch_invariance_slice PASSED [ 30%]
tests/test_rms_norm.py::test_registry_dispatches_rms_norm PASSED [ 31%]
tests/test_rms_norm.py::test_cuda_triton_rms_norm_matches_native_forward_and_backward[1-128-dtype0-triton] SKIPPED [ 33%]
tests/test_rms_norm.py::test_cuda_triton_rms_norm_matches_native_forward_and_backward[1-128-dtype0-cuda] SKIPPED [ 35%]
tests/test_rms_norm.py::test_cuda_triton_rms_norm_matches_native_forward_and_backward[8-768-dtype0-triton] SKIPPED [ 36%]
tests/test_rms_norm.py::test_cuda_triton_rms_norm_matches_native_forward_and_backward[8-768-dtype0-cuda] SKIPPED [ 38%]
tests/test_rms_norm.py::test_cuda_triton_rms_norm_matches_native_forward_and_backward[32-2048-dtype0-triton] SKIPPED [ 40%]
tests/test_rms_norm.py::test_cuda_triton_rms_norm_matches_native_forward_and_backward[32-2048-dtype0-cuda] SKIPPED [ 41%]
tests/test_rms_norm.py::test_cuda_triton_rms_norm_matches_native_forward_and_backward[128-4096-dtype0-triton] SKIPPED [ 43%]
tests/test_rms_norm.py::test_cuda_triton_rms_norm_matches_native_forward_and_backward[128-4096-dtype0-cuda] SKIPPED [ 45%]
tests/test_rms_norm.py::test_cuda_triton_rms_norm_deterministic_repeat[triton] SKIPPED [ 46%]
tests/test_rms_norm.py::test_cuda_triton_rms_norm_deterministic_repeat[cuda] SKIPPED [ 48%]
tests/test_rms_norm.py::test_triton_rms_norm_long_context_dw_reduction SKIPPED [ 50%]
tests/test_rms_norm.py::test_cuda_triton_rms_norm_forward_dx_layout_invariance[63-triton] SKIPPED [ 51%]
tests/test_rms_norm.py::test_cuda_triton_rms_norm_forward_dx_layout_invariance[63-cuda] SKIPPED [ 53%]
tests/test_rms_norm.py::test_cuda_triton_rms_norm_forward_dx_layout_invariance[64-triton] SKIPPED [ 55%]
tests/test_rms_norm.py::test_cuda_triton_rms_norm_forward_dx_layout_invariance[64-cuda] SKIPPED [ 56%]
tests/test_rms_norm.py::test_cuda_triton_rms_norm_forward_dx_layout_invariance[65-triton] SKIPPED [ 58%]
tests/test_rms_norm.py::test_cuda_triton_rms_norm_forward_dx_layout_invariance[65-cuda] SKIPPED [ 60%]
tests/test_rms_norm.py::test_cuda_triton_rms_norm_forward_dx_layout_invariance[127-triton] SKIPPED [ 61%]
tests/test_rms_norm.py::test_cuda_triton_rms_norm_forward_dx_layout_invariance[127-cuda] SKIPPED [ 63%]
tests/test_rms_norm.py::test_cuda_triton_rms_norm_forward_dx_layout_invariance[128-triton] SKIPPED [ 65%]
tests/test_rms_norm.py::test_cuda_triton_rms_norm_forward_dx_layout_invariance[128-cuda] SKIPPED [ 66%]
tests/test_rms_norm.py::test_cuda_triton_rms_norm_forward_dx_layout_invariance[129-triton] SKIPPED [ 68%]
tests/test_rms_norm.py::test_cuda_triton_rms_norm_forward_dx_layout_invariance[129-cuda] SKIPPED [ 70%]
tests/test_rms_norm.py::test_cuda_triton_rms_norm_forward_dx_layout_invariance[255-triton] SKIPPED [ 71%]
tests/test_rms_norm.py::test_cuda_triton_rms_norm_forward_dx_layout_invariance[255-cuda] SKIPPED [ 73%]
tests/test_rms_norm.py::test_cuda_triton_rms_norm_forward_dx_layout_invariance[256-triton] SKIPPED [ 75%]
tests/test_rms_norm.py::test_cuda_triton_rms_norm_forward_dx_layout_invariance[256-cuda] SKIPPED [ 76%]
tests/test_rms_norm.py::test_cuda_triton_rms_norm_forward_dx_layout_invariance[257-triton] SKIPPED [ 78%]
tests/test_rms_norm.py::test_cuda_triton_rms_norm_forward_dx_layout_invariance[257-cuda] SKIPPED [ 80%]
tests/test_rms_norm.py::test_cuda_triton_rms_norm_forward_dx_layout_invariance[4096-triton] SKIPPED [ 81%]
tests/test_rms_norm.py::test_cuda_triton_rms_norm_forward_dx_layout_invariance[4096-cuda] SKIPPED [ 83%]
tests/test_rms_norm.py::test_cuda_rms_norm_masked_dw_layout_invariance SKIPPED [ 85%]
rl_engine/tests/test_dispatch.py::test_logger_enhancements PASSED [ 86%]
rl_engine/tests/test_dispatch.py::test_device_and_registry PASSED [ 88%]
rl_engine/tests/test_dispatch.py::test_rocm_attention_uses_flash_attention_by_default PASSED [ 90%]
rl_engine/tests/test_dispatch.py::test_rocm_attention_native_sdpa_opt_out PASSED [ 91%]
rl_engine/tests/test_dispatch.py::test_registry_explicit_device_selects_device_platform PASSED [ 93%]
rl_engine/tests/test_dispatch.py::test_npu_registry_preserves_cpu_fallbacks_with_ascend_overrides PASSED [ 95%]
rl_engine/tests/test_dispatch.py::test_npu_rmsnorm_falls_back_when_ascend_extension_is_unavailable PASSED [ 96%]
rl_engine/tests/test_dispatch.py::test_npu_available_handles_runtime_failure PASSED [ 98%]
rl_engine/tests/test_dispatch.py::test_executor_flow PASSED [100%]
=============================== warnings summary ===============================
../venv310/lib/python3.10/site-packages/torch/jit/_script.py:362: 14 warnings
/home/ma-user/work/z84450661/venv310/lib/python3.10/site-packages/torch/jit/_script.py:362: DeprecationWarning: `torch.jit.script_method` is deprecated. Please switch to `torch.compile` or `torch.export`.
warnings.warn(
-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
================= 28 passed, 32 skipped, 14 warnings in 4.44s ==================
Here is the result for running tests/test_batch_invariant_logp.py + tests/test_deterministic_logp.py
Details
============================= test session starts ==============================
platform linux -- Python 3.10.14, pytest-9.1.1, pluggy-1.6.0 -- /home/ma-user/work/z84450661/venv310/bin/python
cachedir: .pytest_cache
rootdir: /home/ma-user/work/z84450661/rl-kernel
configfile: pyproject.toml
plugins: anyio-4.14.1
collecting ... collected 160 items
tests/test_batch_invariant_logp.py::TestCorrectness::test_matches_reference_basic PASSED [ 0%]
tests/test_batch_invariant_logp.py::TestCorrectness::test_matches_native_logp_op PASSED [ 1%]
tests/test_batch_invariant_logp.py::TestCorrectness::test_leading_shape_preserved PASSED [ 1%]
tests/test_batch_invariant_logp.py::TestCorrectness::test_bf16_input_fp32_output PASSED [ 2%]
tests/test_batch_invariant_logp.py::TestCorrectness::test_fp16_input_fp32_output PASSED [ 3%]
tests/test_batch_invariant_logp.py::TestCorrectness::test_single_token PASSED [ 3%]
tests/test_batch_invariant_logp.py::TestCorrectness::test_vocab_size_1 PASSED [ 4%]
tests/test_batch_invariant_logp.py::TestCorrectness::test_large_vocab PASSED [ 5%]
tests/test_batch_invariant_logp.py::TestBatchInvariance::test_batch_size_1_vs_n PASSED [ 5%]
tests/test_batch_invariant_logp.py::TestBatchInvariance::test_different_positions_in_batch PASSED [ 6%]
tests/test_batch_invariant_logp.py::TestBatchInvariance::test_mixed_batch_content PASSED [ 6%]
tests/test_batch_invariant_logp.py::TestBatchInvariance::test_padding_layout_invariance PASSED [ 7%]
tests/test_batch_invariant_logp.py::TestBatchInvariance::test_repeated_runs_deterministic PASSED [ 8%]
tests/test_batch_invariant_logp.py::TestBatchInvariance::test_batch_invariance_with_ignore_index PASSED [ 8%]
tests/test_batch_invariant_logp.py::TestValidation::test_rejects_1d_logits PASSED [ 9%]
tests/test_batch_invariant_logp.py::TestValidation::test_rejects_shape_mismatch PASSED [ 10%]
tests/test_batch_invariant_logp.py::TestValidation::test_rejects_negative_target PASSED [ 10%]
tests/test_batch_invariant_logp.py::TestValidation::test_rejects_target_ge_vocab PASSED [ 11%]
tests/test_batch_invariant_logp.py::TestValidation::test_negative_target_with_ignore_index_ok PASSED [ 11%]
tests/test_batch_invariant_logp.py::TestValidation::test_3d_logits PASSED [ 12%]
tests/test_batch_invariant_logp.py::test_ascend_backward_formula_matches_reference_without_full_onehot PASSED [ 13%]
tests/test_batch_invariant_logp.py::test_ascend_backward_formula_accepts_empty_batch PASSED [ 13%]
tests/test_batch_invariant_logp.py::TestBackward::test_backward_matches_reference PASSED [ 14%]
tests/test_batch_invariant_logp.py::TestBackward::test_gradient_batch_invariance PASSED [ 15%]
tests/test_batch_invariant_logp.py::TestIgnoreEdgeCases::test_all_ignore_index_outputs_zero PASSED [ 15%]
tests/test_batch_invariant_logp.py::TestIgnoreEdgeCases::test_custom_ignore_index PASSED [ 16%]
tests/test_batch_invariant_logp.py::TestCUDACorrectness::test_matches_reference_cuda SKIPPED [ 16%]
tests/test_batch_invariant_logp.py::TestCUDACorrectness::test_bf16_cuda SKIPPED [ 17%]
tests/test_batch_invariant_logp.py::TestCUDACorrectness::test_large_vocab_cuda SKIPPED [ 18%]
tests/test_batch_invariant_logp.py::TestCUDABatchInvariance::test_batch_size_1_vs_n_cuda SKIPPED [ 18%]
tests/test_batch_invariant_logp.py::TestCUDABatchInvariance::test_different_positions_cuda SKIPPED [ 19%]
tests/test_batch_invariant_logp.py::TestCUDABatchInvariance::test_repeated_runs_cuda SKIPPED [ 20%]
tests/test_batch_invariant_logp.py::TestCUDABatchInvariance::test_cpu_gpu_cross_check SKIPPED [ 20%]
tests/test_batch_invariant_logp.py::TestTritonCorrectness::test_matches_reference_fp32 SKIPPED [ 21%]
tests/test_batch_invariant_logp.py::TestTritonCorrectness::test_matches_reference_bf16 SKIPPED [ 21%]
tests/test_batch_invariant_logp.py::TestTritonCorrectness::test_matches_reference_fp16 SKIPPED [ 22%]
tests/test_batch_invariant_logp.py::TestTritonCorrectness::test_large_vocab SKIPPED [ 23%]
tests/test_batch_invariant_logp.py::TestTritonCorrectness::test_single_token SKIPPED [ 23%]
tests/test_batch_invariant_logp.py::TestTritonCorrectness::test_3d_logits SKIPPED [ 24%]
tests/test_batch_invariant_logp.py::TestTritonCorrectness::test_matches_pytorch_op SKIPPED [ 25%]
tests/test_batch_invariant_logp.py::TestTritonBatchInvariance::test_batch_size_1_vs_n SKIPPED [ 25%]
tests/test_batch_invariant_logp.py::TestTritonBatchInvariance::test_different_positions SKIPPED [ 26%]
tests/test_batch_invariant_logp.py::TestTritonBatchInvariance::test_repeated_runs SKIPPED [ 26%]
tests/test_batch_invariant_logp.py::TestTritonBatchInvariance::test_mixed_batch_content SKIPPED [ 27%]
tests/test_batch_invariant_logp.py::TestTritonBackward::test_backward_matches_reference SKIPPED [ 28%]
tests/test_batch_invariant_logp.py::TestTritonBackward::test_gradient_batch_invariance SKIPPED [ 28%]
tests/test_batch_invariant_logp.py::TestTritonBackward::test_ignored_row_grad_is_zero SKIPPED [ 29%]
tests/test_batch_invariant_logp.py::TestTritonBackward::test_backward_bf16_input SKIPPED [ 30%]
tests/test_batch_invariant_logp.py::TestTritonBackward::test_backward_fp16_input SKIPPED [ 30%]
tests/test_batch_invariant_logp.py::TestTritonIgnoreIndex::test_ignore_outputs_zero SKIPPED [ 31%]
tests/test_batch_invariant_logp.py::TestTritonIgnoreIndex::test_all_ignore SKIPPED [ 31%]
tests/test_batch_invariant_logp.py::TestTritonCPUValidation::test_rejects_cpu_tensor PASSED [ 32%]
tests/test_batch_invariant_logp.py::TestTritonValidation::test_rejects_1d_logits SKIPPED [ 33%]
tests/test_batch_invariant_logp.py::TestTritonValidation::test_rejects_invalid_target SKIPPED [ 33%]
tests/test_batch_invariant_logp.py::TestCudaSM90Correctness::test_matches_reference_fp32 SKIPPED [ 34%]
tests/test_batch_invariant_logp.py::TestCudaSM90Correctness::test_matches_reference_bf16 SKIPPED [ 35%]
tests/test_batch_invariant_logp.py::TestCudaSM90Correctness::test_large_vocab SKIPPED [ 35%]
tests/test_batch_invariant_logp.py::TestCudaSM90Correctness::test_unaligned_vocab SKIPPED [ 36%]
tests/test_batch_invariant_logp.py::TestCudaSM90Correctness::test_unaligned_row_stride_is_rejected SKIPPED [ 36%]
tests/test_batch_invariant_logp.py::TestCudaSM90Correctness::test_single_token SKIPPED [ 37%]
tests/test_batch_invariant_logp.py::TestCudaSM90Correctness::test_3d_logits SKIPPED [ 38%]
tests/test_batch_invariant_logp.py::TestCudaSM90Correctness::test_matches_pytorch_op SKIPPED [ 38%]
tests/test_batch_invariant_logp.py::TestCudaSM90BatchInvariance::test_batch_size_1_vs_n SKIPPED [ 39%]
tests/test_batch_invariant_logp.py::TestCudaSM90BatchInvariance::test_different_positions SKIPPED [ 40%]
tests/test_batch_invariant_logp.py::TestCudaSM90BatchInvariance::test_repeated_runs SKIPPED [ 40%]
tests/test_batch_invariant_logp.py::TestCudaSM90Backward::test_backward_matches_reference SKIPPED [ 41%]
tests/test_batch_invariant_logp.py::TestCudaSM90Backward::test_ignored_row_grad_is_zero SKIPPED [ 41%]
tests/test_batch_invariant_logp.py::TestCudaSM90IgnoreIndex::test_ignore_outputs_zero SKIPPED [ 42%]
tests/test_batch_invariant_logp.py::TestCudaSM90UnsupportedInputs::test_fp16_is_rejected SKIPPED [ 43%]
tests/test_batch_invariant_logp.py::TestAscendCorrectness::test_matches_reference_fp32 PASSED [ 43%]
tests/test_batch_invariant_logp.py::TestAscendCorrectness::test_matches_reference_bf16 PASSED [ 44%]
tests/test_batch_invariant_logp.py::TestAscendCorrectness::test_large_vocab PASSED [ 45%]
tests/test_batch_invariant_logp.py::TestAscendCorrectness::test_single_token PASSED [ 45%]
tests/test_batch_invariant_logp.py::TestAscendCorrectness::test_3d_logits PASSED [ 46%]
tests/test_batch_invariant_logp.py::TestAscendCorrectness::test_matches_pytorch_op PASSED [ 46%]
tests/test_batch_invariant_logp.py::TestAscendCorrectness::test_empty_batch PASSED [ 47%]
tests/test_batch_invariant_logp.py::TestAscendBatchInvariance::test_batch_size_1_vs_n PASSED [ 48%]
tests/test_batch_invariant_logp.py::TestAscendBatchInvariance::test_different_positions PASSED [ 48%]
tests/test_batch_invariant_logp.py::TestAscendBatchInvariance::test_repeated_runs PASSED [ 49%]
tests/test_batch_invariant_logp.py::TestAscendBackward::test_backward_matches_reference PASSED [ 50%]
tests/test_batch_invariant_logp.py::TestAscendBackward::test_ignored_row_grad_is_zero PASSED [ 50%]
tests/test_batch_invariant_logp.py::TestAscendIgnoreIndex::test_ignore_outputs_zero PASSED [ 51%]
tests/test_batch_invariant_logp.py::TestAscendValidation::test_rejects_invalid_target_when_requested PASSED [ 51%]
tests/test_batch_invariant_logp.py::TestAscendFallback::test_fp16_falls_back PASSED [ 52%]
tests/test_batch_invariant_logp.py::test_registry_dispatches_correctly PASSED [ 53%]
tests/test_batch_invariant_logp.py::test_benchmark_selects_accelerated_op_for_active_device PASSED [ 53%]
tests/test_deterministic_logp.py::test_deterministic_logp_source_locks_reduction_contract PASSED [ 54%]
tests/test_deterministic_logp.py::test_deterministic_logp_shape_dtype_matrix_cuda[fp16-single-token-vocab] SKIPPED [ 55%]
tests/test_deterministic_logp.py::test_deterministic_logp_shape_dtype_matrix_cuda[fp16-tiny-vocab] SKIPPED [ 55%]
tests/test_deterministic_logp.py::test_deterministic_logp_shape_dtype_matrix_cuda[fp16-below-warp] SKIPPED [ 56%]
tests/test_deterministic_logp.py::test_deterministic_logp_shape_dtype_matrix_cuda[fp16-one-warp] SKIPPED [ 56%]
tests/test_deterministic_logp.py::test_deterministic_logp_shape_dtype_matrix_cuda[fp16-above-warp] SKIPPED [ 57%]
tests/test_deterministic_logp.py::test_deterministic_logp_shape_dtype_matrix_cuda[fp16-below-small-bucket] SKIPPED [ 58%]
tests/test_deterministic_logp.py::test_deterministic_logp_shape_dtype_matrix_cuda[fp16-small-bucket-boundary] SKIPPED [ 58%]
tests/test_deterministic_logp.py::test_deterministic_logp_shape_dtype_matrix_cuda[fp16-medium-bucket-start] SKIPPED [ 59%]
tests/test_deterministic_logp.py::test_deterministic_logp_shape_dtype_matrix_cuda[fp16-below-block] SKIPPED [ 60%]
tests/test_deterministic_logp.py::test_deterministic_logp_shape_dtype_matrix_cuda[fp16-one-block] SKIPPED [ 60%]
tests/test_deterministic_logp.py::test_deterministic_logp_shape_dtype_matrix_cuda[fp16-above-block] SKIPPED [ 61%]
tests/test_deterministic_logp.py::test_deterministic_logp_shape_dtype_matrix_cuda[fp16-multi-block-stride] SKIPPED [ 61%]
tests/test_deterministic_logp.py::test_deterministic_logp_shape_dtype_matrix_cuda[fp16-below-medium-boundary] SKIPPED [ 62%]
tests/test_deterministic_logp.py::test_deterministic_logp_shape_dtype_matrix_cuda[fp16-medium-bucket-boundary] SKIPPED [ 63%]
tests/test_deterministic_logp.py::test_deterministic_logp_shape_dtype_matrix_cuda[fp16-large-bucket-start] SKIPPED [ 63%]
tests/test_deterministic_logp.py::test_deterministic_logp_shape_dtype_matrix_cuda[fp16-large-prime-vocab] SKIPPED [ 64%]
tests/test_deterministic_logp.py::test_deterministic_logp_shape_dtype_matrix_cuda[fp16-large-power-two-vocab] SKIPPED [ 65%]
tests/test_deterministic_logp.py::test_deterministic_logp_shape_dtype_matrix_cuda[bf16-single-token-vocab] SKIPPED [ 65%]
tests/test_deterministic_logp.py::test_deterministic_logp_shape_dtype_matrix_cuda[bf16-tiny-vocab] SKIPPED [ 66%]
tests/test_deterministic_logp.py::test_deterministic_logp_shape_dtype_matrix_cuda[bf16-below-warp] SKIPPED [ 66%]
tests/test_deterministic_logp.py::test_deterministic_logp_shape_dtype_matrix_cuda[bf16-one-warp] SKIPPED [ 67%]
tests/test_deterministic_logp.py::test_deterministic_logp_shape_dtype_matrix_cuda[bf16-above-warp] SKIPPED [ 68%]
tests/test_deterministic_logp.py::test_deterministic_logp_shape_dtype_matrix_cuda[bf16-below-small-bucket] SKIPPED [ 68%]
tests/test_deterministic_logp.py::test_deterministic_logp_shape_dtype_matrix_cuda[bf16-small-bucket-boundary] SKIPPED [ 69%]
tests/test_deterministic_logp.py::test_deterministic_logp_shape_dtype_matrix_cuda[bf16-medium-bucket-start] SKIPPED [ 70%]
tests/test_deterministic_logp.py::test_deterministic_logp_shape_dtype_matrix_cuda[bf16-below-block] SKIPPED [ 70%]
tests/test_deterministic_logp.py::test_deterministic_logp_shape_dtype_matrix_cuda[bf16-one-block] SKIPPED [ 71%]
tests/test_deterministic_logp.py::test_deterministic_logp_shape_dtype_matrix_cuda[bf16-above-block] SKIPPED [ 71%]
tests/test_deterministic_logp.py::test_deterministic_logp_shape_dtype_matrix_cuda[bf16-multi-block-stride] SKIPPED [ 72%]
tests/test_deterministic_logp.py::test_deterministic_logp_shape_dtype_matrix_cuda[bf16-below-medium-boundary] SKIPPED [ 73%]
tests/test_deterministic_logp.py::test_deterministic_logp_shape_dtype_matrix_cuda[bf16-medium-bucket-boundary] SKIPPED [ 73%]
tests/test_deterministic_logp.py::test_deterministic_logp_shape_dtype_matrix_cuda[bf16-large-bucket-start] SKIPPED [ 74%]
tests/test_deterministic_logp.py::test_deterministic_logp_shape_dtype_matrix_cuda[bf16-large-prime-vocab] SKIPPED [ 75%]
tests/test_deterministic_logp.py::test_deterministic_logp_shape_dtype_matrix_cuda[bf16-large-power-two-vocab] SKIPPED [ 75%]
tests/test_deterministic_logp.py::test_deterministic_logp_shape_dtype_matrix_cuda[fp32-single-token-vocab] SKIPPED [ 76%]
tests/test_deterministic_logp.py::test_deterministic_logp_shape_dtype_matrix_cuda[fp32-tiny-vocab] SKIPPED [ 76%]
tests/test_deterministic_logp.py::test_deterministic_logp_shape_dtype_matrix_cuda[fp32-below-warp] SKIPPED [ 77%]
tests/test_deterministic_logp.py::test_deterministic_logp_shape_dtype_matrix_cuda[fp32-one-warp] SKIPPED [ 78%]
tests/test_deterministic_logp.py::test_deterministic_logp_shape_dtype_matrix_cuda[fp32-above-warp] SKIPPED [ 78%]
tests/test_deterministic_logp.py::test_deterministic_logp_shape_dtype_matrix_cuda[fp32-below-small-bucket] SKIPPED [ 79%]
tests/test_deterministic_logp.py::test_deterministic_logp_shape_dtype_matrix_cuda[fp32-small-bucket-boundary] SKIPPED [ 80%]
tests/test_deterministic_logp.py::test_deterministic_logp_shape_dtype_matrix_cuda[fp32-medium-bucket-start] SKIPPED [ 80%]
tests/test_deterministic_logp.py::test_deterministic_logp_shape_dtype_matrix_cuda[fp32-below-block] SKIPPED [ 81%]
tests/test_deterministic_logp.py::test_deterministic_logp_shape_dtype_matrix_cuda[fp32-one-block] SKIPPED [ 81%]
tests/test_deterministic_logp.py::test_deterministic_logp_shape_dtype_matrix_cuda[fp32-above-block] SKIPPED [ 82%]
tests/test_deterministic_logp.py::test_deterministic_logp_shape_dtype_matrix_cuda[fp32-multi-block-stride] SKIPPED [ 83%]
tests/test_deterministic_logp.py::test_deterministic_logp_shape_dtype_matrix_cuda[fp32-below-medium-boundary] SKIPPED [ 83%]
tests/test_deterministic_logp.py::test_deterministic_logp_shape_dtype_matrix_cuda[fp32-medium-bucket-boundary] SKIPPED [ 84%]
tests/test_deterministic_logp.py::test_deterministic_logp_shape_dtype_matrix_cuda[fp32-large-bucket-start] SKIPPED [ 85%]
tests/test_deterministic_logp.py::test_deterministic_logp_shape_dtype_matrix_cuda[fp32-large-prime-vocab] SKIPPED [ 85%]
tests/test_deterministic_logp.py::test_deterministic_logp_shape_dtype_matrix_cuda[fp32-large-power-two-vocab] SKIPPED [ 86%]
tests/test_deterministic_logp.py::test_deterministic_logp_repeatability_cuda SKIPPED [ 86%]
tests/test_deterministic_logp.py::test_deterministic_logp_out_dtype_matrix_reuses_storage_cuda[fp16] SKIPPED [ 87%]
tests/test_deterministic_logp.py::test_deterministic_logp_out_dtype_matrix_reuses_storage_cuda[bf16] SKIPPED [ 88%]
tests/test_deterministic_logp.py::test_deterministic_logp_out_dtype_matrix_reuses_storage_cuda[fp32] SKIPPED [ 88%]
tests/test_deterministic_logp.py::test_deterministic_logp_out_dtype_matrix_reuses_storage_cuda[fp64] SKIPPED [ 89%]
tests/test_deterministic_logp.py::test_deterministic_logp_non_contiguous_inputs_cuda SKIPPED [ 90%]
tests/test_deterministic_logp.py::test_deterministic_logp_batch_size_invariance_cuda SKIPPED [ 90%]
tests/test_deterministic_logp.py::test_deterministic_logp_batch_position_invariance_cuda SKIPPED [ 91%]
tests/test_deterministic_logp.py::test_deterministic_logp_bucket_boundaries_are_batch_and_indexed_invariant_cuda[128] SKIPPED [ 91%]
tests/test_deterministic_logp.py::test_deterministic_logp_bucket_boundaries_are_batch_and_indexed_invariant_cuda[129] SKIPPED [ 92%]
tests/test_deterministic_logp.py::test_deterministic_logp_bucket_boundaries_are_batch_and_indexed_invariant_cuda[4096] SKIPPED [ 93%]
tests/test_deterministic_logp.py::test_deterministic_logp_bucket_boundaries_are_batch_and_indexed_invariant_cuda[4097] SKIPPED [ 93%]
tests/test_deterministic_logp.py::test_deterministic_logp_ignores_batch_noise_bitwise_cuda SKIPPED [ 94%]
tests/test_deterministic_logp.py::test_deterministic_logp_indexed_matches_dense_bits_cuda SKIPPED [ 95%]
tests/test_deterministic_logp.py::test_deterministic_logp_indexed_out_preserves_inactive_rows_cuda SKIPPED [ 95%]
tests/test_deterministic_logp.py::test_deterministic_logp_indexed_fp32_empty_indices_zero_fills_cuda SKIPPED [ 96%]
tests/test_deterministic_logp.py::test_deterministic_logp_invalid_token_ids_zero_fill_cuda SKIPPED [ 96%]
tests/test_deterministic_logp.py::test_deterministic_logp_extreme_logits_are_stable_cuda SKIPPED [ 97%]
tests/test_deterministic_logp.py::test_deterministic_logp_rejects_bad_shapes_and_output_dtype_cuda SKIPPED [ 98%]
tests/test_deterministic_logp.py::test_deterministic_logp_out_of_range_indices_do_not_overwrite_output_cuda SKIPPED [ 98%]
tests/test_deterministic_logp.py::test_deterministic_logp_matches_reference_tolerance_cuda[dtype0] SKIPPED [ 99%]
tests/test_deterministic_logp.py::test_deterministic_logp_matches_reference_tolerance_cuda[dtype1] SKIPPED [100%]
=============================== warnings summary ===============================
../venv310/lib/python3.10/site-packages/torch/jit/_script.py:362: 14 warnings
/home/ma-user/work/z84450661/venv310/lib/python3.10/site-packages/torch/jit/_script.py:362: DeprecationWarning: `torch.jit.script_method` is deprecated. Please switch to `torch.compile` or `torch.export`.
warnings.warn(
tests/test_batch_invariant_logp.py::TestAscendCorrectness::test_matches_pytorch_op
/home/ma-user/work/z84450661/rl-kernel/rl_engine/kernels/ops/pytorch/loss/batch_invariant_logp.py:105: UserWarning: Cannot create tensor with interal format while allow_internel_format=False, tensor will be created with base format. (Triggered internally at ../torch_npu/csrc/aten/common/TensorFactories.cpp:340.)
selected_logp = selected_logp.where(valid_mask, torch.zeros_like(selected_logp))
-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
================= 45 passed, 115 skipped, 15 warnings in 5.85s =================
Here is all the gtest results, which passes, scripts/check_operator.py --check-grad
Details
5. rms_norm / fp32
INFO 08-21 15:13:48 [RL-Kernel]: Successfully linked to precompiled _C_npu.rmsnorm_ascend_forward kernel.
/home/ma-user/work/z84450661/venv310/lib/python3.10/site-packages/torch/autograd/function.py:583: UserWarning: Cannot create tensor with interal format while allow_internel_format=False, tensor will be created with base format. (Triggered internally at ../torch_npu/csrc/aten/common/TensorFactories.cpp:340.)
return super().apply(*args, **kwargs) # type: ignore[misc]
suite=rms_norm passed=True pass_rate=1.0000
candidate=ascend-rms_norm backend=ascend passed=True pass_rate=1.0000
case=rms_norm-torch.float32-2x16x4096 output=0 shape=(2, 16, 4096) dtype=torch.float32 max_abs=4.48226929e-05 mean_abs=1.27800877e-06 max_rel=5.83826613e-06 tol=(atol=1.000e-04, rtol=1.000e-04) passed=True
case=rms_norm-torch.float32-2x16x4096 output=1 gradient:x shape=(2, 16, 4096) dtype=torch.float32 max_abs=9.53674316e-07 mean_abs=1.82109936e-08 max_rel=2.08102501e-04 tol=(atol=1.000e-04, rtol=1.000e-04) passed=True
case=rms_norm-torch.float32-2x16x4096 output=2 gradient:weight shape=(4096,) dtype=torch.float32 max_abs=1.52587891e-05 mean_abs=2.26008706e-06 max_rel=3.45694730e-07 tol=(atol=1.000e-04, rtol=1.000e-04) passed=True
5. qk_norm / fp32
INFO 08-21 15:13:57 [RL-Kernel]: Successfully linked to precompiled _C_npu.rmsnorm_ascend_forward kernel.
/home/ma-user/work/z84450661/venv310/lib/python3.10/site-packages/torch/autograd/function.py:583: UserWarning: Cannot create tensor with interal format while allow_internel_format=False, tensor will be created with base format. (Triggered internally at ../torch_npu/csrc/aten/common/TensorFactories.cpp:340.)
return super().apply(*args, **kwargs) # type: ignore[misc]
suite=qk_norm passed=True pass_rate=1.0000
candidate=ascend-qk_norm backend=ascend passed=True pass_rate=1.0000
case=qk_norm-torch.float32-2x16x32x128 output=0 shape=(2, 512, 128) dtype=torch.float32 max_abs=6.62803650e-05 mean_abs=1.36790311e-06 max_rel=1.25606584e-05 tol=(atol=1.000e-04, rtol=1.000e-04) passed=True
case=qk_norm-torch.float32-2x16x32x128 output=1 gradient:x shape=(2, 512, 128) dtype=torch.float32 max_abs=9.53674316e-07 mean_abs=1.87886435e-08 max_rel=6.20549684e-03 tol=(atol=1.000e-04, rtol=1.000e-04) passed=True
case=qk_norm-torch.float32-2x16x32x128 output=2 gradient:weight shape=(128,) dtype=torch.float32 max_abs=1.34277344e-03 mean_abs=4.03404236e-04 max_rel=1.32250932e-06 tol=(atol=1.000e-04, rtol=1.000e-04) passed=True
5. rms_norm / bf16
INFO 08-21 15:14:06 [RL-Kernel]: Successfully linked to precompiled _C_npu.rmsnorm_ascend_forward kernel.
/home/ma-user/work/z84450661/venv310/lib/python3.10/site-packages/torch/autograd/function.py:583: UserWarning: Cannot create tensor with interal format while allow_internel_format=False, tensor will be created with base format. (Triggered internally at ../torch_npu/csrc/aten/common/TensorFactories.cpp:340.)
return super().apply(*args, **kwargs) # type: ignore[misc]
suite=rms_norm passed=True pass_rate=1.0000
candidate=ascend-rms_norm backend=ascend passed=True pass_rate=1.0000
case=rms_norm-torch.bfloat16-2x16x4096 output=0 shape=(2, 16, 4096) dtype=torch.bfloat16 max_abs=2.73237228e-02 mean_abs=8.88527837e-04 max_rel=3.88927641e-03 tol=(atol=5.000e-02, rtol=2.000e-02) passed=True
case=rms_norm-torch.bfloat16-2x16x4096 output=1 gradient:x shape=(2, 16, 4096) dtype=torch.bfloat16 max_abs=7.45058060e-09 mean_abs=5.68434189e-14 max_rel=4.34782589e-03 tol=(atol=1.000e-01, rtol=2.000e-02) passed=True
case=rms_norm-torch.bfloat16-2x16x4096 output=2 gradient:weight shape=(4096,) dtype=torch.bfloat16 max_abs=0.00000000e+00 mean_abs=0.00000000e+00 max_rel=0.00000000e+00 tol=(atol=1.000e-01, rtol=2.000e-02) passed=True
5. qk_norm / bf16
INFO 08-21 15:14:15 [RL-Kernel]: Successfully linked to precompiled _C_npu.rmsnorm_ascend_forward kernel.
/home/ma-user/work/z84450661/venv310/lib/python3.10/site-packages/torch/autograd/function.py:583: UserWarning: Cannot create tensor with interal format while allow_internel_format=False, tensor will be created with base format. (Triggered internally at ../torch_npu/csrc/aten/common/TensorFactories.cpp:340.)
return super().apply(*args, **kwargs) # type: ignore[misc]
suite=qk_norm passed=True pass_rate=1.0000
candidate=ascend-qk_norm backend=ascend passed=True pass_rate=1.0000
case=qk_norm-torch.bfloat16-2x16x32x128 output=0 shape=(2, 512, 128) dtype=torch.bfloat16 max_abs=2.87456512e-02 mean_abs=8.52997880e-04 max_rel=3.89306131e-03 tol=(atol=5.000e-02, rtol=2.000e-02) passed=True
case=qk_norm-torch.bfloat16-2x16x32x128 output=1 gradient:x shape=(2, 512, 128) dtype=torch.bfloat16 max_abs=3.81469727e-06 mean_abs=2.91038305e-11 max_rel=4.46428591e-03 tol=(atol=1.000e-01, rtol=2.000e-02) passed=True
case=qk_norm-torch.bfloat16-2x16x32x128 output=2 gradient:weight shape=(128,) dtype=torch.bfloat16 max_abs=0.00000000e+00 mean_abs=0.00000000e+00 max_rel=0.00000000e+00 tol=(atol=1.000e-01, rtol=2.000e-02) passed=True
5. rms_norm / fp16
INFO 08-21 15:14:24 [RL-Kernel]: Successfully linked to precompiled _C_npu.rmsnorm_ascend_forward kernel.
/home/ma-user/work/z84450661/venv310/lib/python3.10/site-packages/torch/autograd/function.py:583: UserWarning: Cannot create tensor with interal format while allow_internel_format=False, tensor will be created with base format. (Triggered internally at ../torch_npu/csrc/aten/common/TensorFactories.cpp:340.)
return super().apply(*args, **kwargs) # type: ignore[misc]
suite=rms_norm passed=True pass_rate=1.0000
candidate=ascend-rms_norm backend=ascend passed=True pass_rate=1.0000
case=rms_norm-torch.float16-2x16x4096 output=0 shape=(2, 16, 4096) dtype=torch.float16 max_abs=3.42178345e-03 mean_abs=1.10994937e-04 max_rel=3.81875038e-03 tol=(atol=1.000e-03, rtol=1.000e-03) passed=True
case=rms_norm-torch.float16-2x16x4096 output=1 gradient:x shape=(2, 16, 4096) dtype=torch.float16 max_abs=1.95312500e-03 mean_abs=3.98285920e-08 max_rel=9.19963175e-04 tol=(atol=1.000e-03, rtol=1.000e-03) passed=True
case=rms_norm-torch.float16-2x16x4096 output=2 gradient:weight shape=(4096,) dtype=torch.float16 max_abs=0.00000000e+00 mean_abs=0.00000000e+00 max_rel=0.00000000e+00 tol=(atol=1.000e-03, rtol=1.000e-03) passed=True
5. qk_norm / fp16
INFO 08-21 15:14:33 [RL-Kernel]: Successfully linked to precompiled _C_npu.rmsnorm_ascend_forward kernel.
/home/ma-user/work/z84450661/venv310/lib/python3.10/site-packages/torch/autograd/function.py:583: UserWarning: Cannot create tensor with interal format while allow_internel_format=False, tensor will be created with base format. (Triggered internally at ../torch_npu/csrc/aten/common/TensorFactories.cpp:340.)
return super().apply(*args, **kwargs) # type: ignore[misc]
suite=qk_norm passed=True pass_rate=1.0000
candidate=ascend-qk_norm backend=ascend passed=True pass_rate=1.0000
case=qk_norm-torch.float16-2x16x32x128 output=0 shape=(2, 512, 128) dtype=torch.float16 max_abs=3.57532501e-03 mean_abs=1.06467298e-04 max_rel=4.61618090e-03 tol=(atol=1.000e-03, rtol=1.000e-03) passed=True
case=qk_norm-torch.float16-2x16x32x128 output=1 gradient:x shape=(2, 512, 128) dtype=torch.float16 max_abs=1.95312500e-03 mean_abs=3.38222890e-08 max_rel=9.38086305e-04 tol=(atol=1.000e-03, rtol=1.000e-03) passed=True
case=qk_norm-torch.float16-2x16x32x128 output=2 gradient:weight shape=(128,) dtype=torch.float16 max_abs=0.00000000e+00 mean_abs=0.00000000e+00 max_rel=0.00000000e+00 tol=(atol=1.000e-03, rtol=1.000e-03) passed=True
…ith-slack docs: replace expired Discord links with Slack
erfgss
force-pushed
the
feat/ascend-rmsnorm
branch
from
August 23, 2026 23:38
d30ed50 to
ca980ae
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.