feat(profiler): support Ascend NPU profiling via torch_npu.profiler#1093
Open
brandneway wants to merge 2 commits into
Open
feat(profiler): support Ascend NPU profiling via torch_npu.profiler#1093brandneway wants to merge 2 commits into
brandneway wants to merge 2 commits into
Conversation
added 2 commits
July 24, 2026 02:08
Switch the NPU path from export_chrome_trace (single json) to on_trace_ready=tensorboard_trace_handler + experimental_config, so the full Ascend profiling output directory is produced: ASCEND_PROFILER_OUTPUT (kernel_details.csv, step_trace_time.csv, trace_view.json, op_summary.csv, op_statistic.csv, operator_details.csv, .db) plus PROF_XXX, FRAMEWORK and logs. data_simplification=False keeps every generated file. - profiler_level configurable via CACHE_DIT_NPU_PROFILER_LEVEL (Level0 default) - sanitize worker_name to the allowed [A-Za-z0-9_-] charset - locate the produced *_ascend_pt dir and expose it via trace_path - CUDA path unchanged
Member
|
Thanks for you contribution! By the way, can you should some test logs about this feature? |
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
2 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.
Summary
cache_dit.generate ... --profilecrashed immediately on Ascend NPU becauseProfilerContexthard-asserteddevice_type == "cuda". This PR adds NPU profiling support so--profileworks end-to-end on Ascend devices.Root cause:
src/cache_dit/profiler.pyasserted CUDA-only in both__init__and__enter__, and used CUDA-specific APIs (torch.profiler.ProfilerActivity.CUDA,torch.cuda.synchronize(),torch.cuda.memory._record_memory_history).Changes (
src/cache_dit/profiler.py)torch_npu.profiler(which mirrors thetorch.profilerAPI); keeptorch.profileron CUDA. Lazy-importtorch_npuonly when on NPU — zero impact on CUDA/CPU.cudaornpu(clearer error message otherwise). CPU remains unsupported."GPU"maps to the platform accelerator activity (CUDA on CUDA, NPU on NPU); the accelerator name ("NPU") is also accepted as a synonym. Default activities["CPU","GPU"]just work on NPU.current_platform.synchronize()instead oftorch.cuda.synchronize().profile_memory=True(whenMEMrequested); the CUDA-onlytorch.cuda.memoryhistory/snapshot path is preserved unchanged for CUDA.torch_npu.profiler.export_chrome_traceonly writes a plain.json, so on NPU the trace is exported to a temp json and gzip-compressed to keep the.trace.json.gzartifact convention used on CUDA.Tests
tests/utils/test_profiler.py(CPU-safe): covers_supported_device(),_resolve_profiler_backend()shape, and the unsupported-device assertion guard. Passes on CPU CI (device-required case is skipped).Verification (Ascend 910B4, container cache-dit-gty)
--profileraisedAssertionError: Torch ProfilerContext currently only supports CUDA devices.python -m cache_dit.generate wan2.2_t2v --model-path ... --cpu-offload --attn _mindiesd_laser --compile --profilenow enters/exits the profiler, collects CANN data, and exports a valid chrome trace (.trace.json.gz, parseable json).Notes
current_platform.device_type.--profile-with-stack+--profile-record-shapes+ 30 steps), NPU traces can be large andtorch_npuexport parsing is the dominant cost — inherent to the trace content, not this change.