UC-01: Extract _softmax_dropout_reference + _prepare_input helpers (Ascend NPU adaptation)

[slamdunk111]

Summary

Extract the generic PyTorch fallback path from softmax_dropout into two standalone, independently testable helpers:

• _prepare_input(input, inplace): single preprocessing point (contiguous + conditional clone)
• _softmax_dropout_reference(...) / _softmax_dropout_reference_prepared(...): generic F.dropout(F.softmax(...)) path

This is a pure extraction with zero behavior change. It fixes the duplicate-clone regression where input preprocessing happened in multiple places.

Ascend NPU Adaptation

This PR is the foundational refactoring for Ascend NPU adaptation of Uni-Core's softmax_dropout module, targeting Ascend NPUs such as 910B2, 910C, and future variants.

On Ascend NPUs (910B2, 910C, etc.), the CUDA fused path (SoftmaxDropoutFast) is unavailable — softmax_dropout always falls back to the generic F.dropout(F.softmax(...)) reference path. Before this PR, that reference path was buried inline inside the public function, untestable in isolation, and could not serve as a verified correctness baseline for Ascend NPU.

What this PR does for Ascend:

1. Extracts the Ascend NPU computation path as a standalone, testable function — _softmax_dropout_reference is the exact code path that runs on Ascend NPUs (910B2/910C). It can now be unit-tested without CUDA hardware.
2. Eliminates duplicate clone — on Ascend NPU, unnecessary tensor copies waste device memory. _prepare_input is called exactly once, reducing Ascend NPU memory pressure.
3. Provides correctness baseline for future Ascend optimization — any future Ascend-specific fused operator can be A/B tested against this verified reference path.
4. Verified end-to-end on Ascend 910B2 (as the available representative of the Ascend NPU family, which also includes 910C and future variants) — 3-seed QM9 A/B training comparison confirms zero regression (avg relative delta 0.88%).

Changes

• unicore/modules/softmax_dropout.py: extract helpers, public softmax_dropout calls _prepare_input exactly once

Why

On non-CUDA backends such as Ascend NPU (Ascend 910B2, 910C, and future variants, using torch_npu), softmax_dropout falls back to a generic F.dropout(F.softmax(...)) path. Extracting this into a standalone helper makes it independently testable without CUDA hardware, and provides a verified baseline for Ascend NPU adaptation.

Quality Assurance

This PR is part of a backend-neutral refactor series (UC-01~UC-03) developed for Ascend NPU adaptation. Full QA artifacts are in the patch Mirror repo.

Test Results (2026-07-11, Ascend 910B2 NPU, torch_npu 2.7.1.post2, Python 3.11, PyTorch 2.7.1)

Test Class	Tests	Result
TestReferenceNumerical	16	ALL PASS
TestPublicApi	7	ALL PASS
TestCheckMask	8	ALL PASS
TestCheckBias	6	ALL PASS
TestDispatch	5	ALL PASS
TestCudaFusedAvailable	8	ALL PASS
Total	50	0 failures, 0 errors, 0 skipped

Ascend NPU E2E Verification

Gate	Environment	Result
E2E-6 AccuracyProxy	Ascend 910B2 NPU	✅ PASS: 3-seed QM9 A/B, avg rel delta 0.88%
E2E-7 FullAccuracy	Ascend 910B2 NPU	✅ PASS: 6 runs (3 seeds × 2 versions × 3 epochs), all <5% threshold
E2E-2 NPU UT parity	Ascend 910B2 NPU	✅ PASS: 27 tests
E2E-4 Dropout semantics	Ascend 910B2 NPU	✅ PASS: 7 tests (1 skipped: torch_npu 2.7.1.post2 limitation)

Coverage

Name                                                         Stmts   Miss Branch BrPart  Cover
PR-1_backend_neutral/unicore_softmax_dropout_refactored.py      91     13     46      7    82%

82% is the non-CUDA-reachable ceiling — CUDA-only code paths (SoftmaxDropoutFast, _cuda_softmax_dropout) are excluded via exclude_lines as they require a real CUDA device.

Key Verified Properties

1. No duplicate clone: _prepare_input called exactly once in public softmax_dropout(). Reference helper receives already-prepared input. Reduces Ascend NPU memory allocations.
2. Mask/bias via add_: Since input is a safe copy (cloned by _prepare_input when inplace=False), mask/bias applied in-place — matches upstream allocation semantics, no extra tensor allocations on Ascend.
3. API signature unchanged: softmax_dropout(input, dropout_prob, is_training, mask, bias, inplace) — identical to upstream.
4. inplace semantics: inplace=True modifies input; inplace=False preserves input.
5. dtype support: fp32, fp16, bf16 all produce correct softmax output — verified on Ascend 910B2 (representative of Ascend NPU family including 910C).
6. 4D input: [bsz, n_heads, tgt, src] shape works correctly on Ascend NPU.

CI Configuration

A 5-job CI workflow is configured (.github/workflows/ci.yml):

1. static-and-security: token leak scan + py_compile all files
2. cpu-characterization: pytest + coverage on Python 3.10/3.11
3. patch-series-verify: git am all patches + cmp against refactored source
4. cuda-regression: placeholder (requires self-hosted CUDA runner)
5. npu-draft-tests: placeholder (requires self-hosted Ascend runner)

Test Plan

• Characterization tests lock current behavior (must stay green on Ascend NPU)
• Unit tests for _prepare_input and _softmax_dropout_reference
• Public API signature unchanged
• Patch applies cleanly via git am and produces refactored source
• Ascend 910B2 E2E: 3-seed QM9 A/B training, all non-inferior
• Independent CI verification

Zero behavior change: the refactored code produces byte-identical output to the upstream else branch for all tested inputs (fp32, fp16, bf16; 3D and 4D; with/without mask/bias) — verified on Ascend 910B2 NPU (representative of the Ascend NPU family, which also includes 910C and future variants).

Patch Source

This PR is generated from a patch mirror at cnpc-chem-opt. The patch file 0001-extract-reference-helper.patch applies cleanly via git am. See also the umbrella issue #76 for the full Ascend NPU adaptation overview.