Skip to content

[NPU][autotuner] config-layer tolerance for Ascend NPU - #3234

Open
yzchang-plus wants to merge 1 commit into
pytorch:mainfrom
yzchang-plus:npu-pr1-config
Open

[NPU][autotuner] config-layer tolerance for Ascend NPU#3234
yzchang-plus wants to merge 1 commit into
pytorch:mainfrom
yzchang-plus:npu-pr1-config

Conversation

@yzchang-plus

Copy link
Copy Markdown

Makes Helion's autotuner/config layer robust on Ascend NPU (UB=192KB, triton-ascend) so NPU-incompatible config values degrade silently instead of crashing, and overflowing configs are skipped rather than aborting autotune.

Changes:

  • block_id_sequence.py: _normalize discards user values when size==0.
  • config_spec.py: downgrade unsupported indexing to NPU ("pointer"); pid_type defaults to "flat"; NPU reduction cap (prod(block)reduction <= UB budget); [None]->default loop; _shrink_for_numel_constraints; imports npu caps.
  • logger.py: classify "ub overflow"/"ConvertLinalgRToBinary" as debug.
  • local_cache.py: npu_volatile_triton_cache_dir (avoid stale binaries).
  • benchmark_provider.py: NPU baseline fallback (halve block_sizes on overflow).
  • benchmarking.py / progress_bar.py: do_bench_npu routing + NPU progress.
  • _compat.py: safe_clear_cache (clear triton cache for NPU drivers without clear_cache); also holds register_npu_backend / supports_launch_cooperative_grid used by later PRs.
  • compile_environment.py: numel constraint substitutes specialized vars and fixed block symbols (so [chunk=256, block_v]-style constraints are created).
  • ascend/config.py (+ascend/init.py): npu* env-tunable caps, in subdir.

Problems solved: autotune crashes on NPU-incompatible config keys; UB-overflow configs aborting autotune; stale cached binaries; baseline overflow; numel constraint missing fixed-block dims.

Test Results

NPU (Ascend 910B4)

This PR is the config/autotuner infrastructure layer. It does not run NPU kernels standalone, but is verified as part of the full 4-PR stack (see PR4). On Ascend910B4 with HELION_BACKEND=ascend, the full stack passes add / matmul / softmax / rms_norm / gdn_fwd_h / se_block EXIT=0.

NPU reproduction

HELION_BACKEND=ascend  \
HELION_NPU_MAX_TENSOR_NUMEL=4096 HELION_NPU_UB_BUDGET_ELEMENTS=1024 \
python examples/<op>.py

PR stack

This is PR 1 of a 4-PR stack (merge in order):

  1. [NPU][autotuner] config-layer tolerance (this PR)
  2. [NPU][codegen] Ascend-safe codegen hooks
  3. [NPU][backend] AscendBackend + AST transforms
  4. [NPU][runtime] runtime launcher + testing + examples

Makes Helion's autotuner/config layer robust on Ascend NPU (UB=192KB,
triton-ascend) so NPU-incompatible config values degrade silently instead of
crashing, and overflowing configs are skipped rather than aborting autotune.

Changes:
- block_id_sequence.py: _normalize discards user values when size==0.
- config_spec.py: downgrade unsupported indexing to NPU ("pointer"); pid_type
  defaults to "flat"; NPU reduction cap (prod(block)*reduction <= UB budget);
  [None]->default loop; _shrink_for_numel_constraints; imports _npu_* caps.
- logger.py: classify "ub overflow"/"ConvertLinalgRToBinary" as debug.
- local_cache.py: npu_volatile_triton_cache_dir (avoid stale binaries).
- benchmark_provider.py: NPU baseline fallback (halve block_sizes on overflow).
- benchmarking.py / progress_bar.py: do_bench_npu routing + NPU progress.
- _compat.py: safe_clear_cache (clear triton cache for NPU drivers without clear_cache); also holds register_npu_backend / supports_launch_cooperative_grid used by later PRs.
- compile_environment.py: numel constraint substitutes specialized vars and
  fixed block symbols (so [chunk=256, block_v]-style constraints are created).
- ascend/config.py (+ascend/__init__.py): _npu_* env-tunable caps, in subdir.

Problems solved: autotune crashes on NPU-incompatible config keys; UB-overflow
configs aborting autotune; stale cached binaries; baseline overflow; numel
constraint missing fixed-block dims.

Test report: config-layer PR; compiles/imports on origin/main. Verified in the
full NPU chain: add/matmul/softmax/rms_norm/gdn EXIT=0 on Ascend910B4.
NPU-only guards; CUDA unaffected.
@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Meta Open Source bot. label Aug 3, 2026
@yzchang-plus

Copy link
Copy Markdown
Author

Hi @jansel
I noticed this PR hasn't been reviewed yet. I understand you're probably busy — just wanted to check if there's anything I can do on my side to make it easier to review. Thanks for your time!

@yzchang-plus

Copy link
Copy Markdown
Author

Hi @jansel I noticed this PR hasn't been reviewed yet. I understand you're probably busy — just wanted to check if there's anything I can do on my side to make it easier to review. Thanks for your time!

this PR is a follow-up to #3037, this PR makes the autotuner/config layer robust on Ascend NPU — incompatible configs degrade silently instead of crashing autotune, and UB-overflow configs are skipped rather than aborting. Changes are strictly NPU-only guards. Verified end-to-end on Ascend910B4. Happy to walk through or split if needed.

@jansel

jansel commented Aug 7, 2026

Copy link
Copy Markdown
Contributor

@yzchang-plus how are we testing this? Are all tests now passing if you run:

HELION_BACKEND=ascend pytest test

Looks like you have merge conflicts.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Meta Open Source bot.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants