feat: add complete MetaX profiler tracing

[lvyufeng]

AI Agent Information

• Agent/Tool: Claude Code CLI
• Model: Claude Opus 5
• Human Reviewer: @aoyulong
• Session Summary: Implemented complete MetaX MCPTI profiler support after comparing the existing CUDA and DCU tracer architecture. The implementation was measured on a C550 with MACA 3.8.0 and rebased onto the latest flagos/main before submission.

Summary

This change adds complete profiler activity collection for MetaX through the CUDA-compatible MCPTI API. It emits kernel, memcpy, memset, and runtime events with correct callback names, correlation flows, device-time attribution, kernel metadata, and capture-window filtering. It also adds per-accelerator tracer selection with an explicit unavailable fallback and documents the measured MetaX support and remaining validation limits.

Change Type

• Bug Fix
• New Feature
• Performance Optimization
• Refactoring
• Documentation
• Testing
• CI/Infrastructure
• Breaking Change

Platforms Affected

• CUDA
• MetaX
• Ascend (source-selection compatibility only; no runtime behavior change)
• PPU (uses the existing CUDA-compatible source-selection path)
• Platform-agnostic (tracer fallback and interface selection)

Problem Analysis

What was broken/missing?

MetaX profiling did not provide a complete device timeline through the existing CUPTI-compatible tracer. Runtime callback IDs were interpreted with the NVIDIA namespace, device-side activity records could be lost or discarded, occupancy metadata was unavailable, and unsupported platforms had no explicit device-tracer factory fallback.

Why did it happen?

MetaX MCPTI uses a different callback-ID namespace and vendor record layouts from NVIDIA CUPTI. MCPTI 3.8.0 can also return a stale iterator pointer after external-correlation records. Its compatibility runtime exposes wcudaDeviceGetAttribute rather than the NVIDIA symbol used by the occupancy calculator. Finally, the source tree uses recursive globbing, so every build must exclude all tracer implementations except exactly one factory.

Investigation process:

1. Compared the generic Kineto adaptor, CUDA/CUPTI tracer, ROCtracer implementation, CMake source selection, and existing profiler parity tests.
2. Probed the MACA 3.8.0 MCPTI library and vendor headers to verify exported symbols, callback IDs, record sizes, activity layouts, and the wcudaDeviceGetAttribute entry point.
3. Reproduced MCPTI resolver deadlocks, repeated iterator pointers, malformed buffer advancement, and zero-timestamp capture-window failures on the C550 workload.
4. Rebuilt MetaX in boxing mode and validated the full parity workload after each parser and build-system correction.

Solution Design

Implementation approach:

• Extend the dlopen shim with MetaX-specific callback ABI types, MCPTI library selection, callback-ID fallback names, and deferred mcptiActivityGetApiName resolution.
• Add a MetaX record scanner that validates aligned candidates, advances by complete vendor record sizes, resynchronizes after malformed data, and skips documented no-timestamp sentinels.
• Resolve MetaX occupancy attributes through wcudaDeviceGetAttribute while retaining the existing CUDA path for NVIDIA and PPU-compatible builds.
• Select exactly one tracer implementation per accelerator and provide an explicit unavailable tracer for platforms without a supported activity API.
• Extend tests and documentation with MetaX-specific loading and measured parity coverage.

Key design decisions:

• Keep vendor-specific activity record mirrors and parsing below the vendor-agnostic DeviceTracer interface.
• Preserve vendor correlation IDs for runtime-to-device flow arrows and PyTorch external correlation IDs for CPU-op device-time attribution; these namespaces are not interchangeable.
• Defer the MetaX API-name resolver until after activity flushing because invoking it from the MCPTI completion callback can deadlock.
• Treat MetaX scanner recovery as a compatibility workaround for the observed MACA 3.8.0 iterator behavior, and document that it still needs validation across additional SDK versions and devices.

Code changes by file:

• csrc/CMakeLists.txt: select exactly one profiler tracer per accelerator, enable MetaX MCPTI definitions, and search the MetaX occupancy compatibility include path.
• csrc/profiler/cupti_shim.h: add MetaX callback ABI declarations, MCPTI library loading, callback-ID mapping, and deferred resolver support.
• csrc/profiler/cupti_device_tracer.cc: parse MCPTI records, recover malformed iterator sequences, skip missing timestamps, resolve occupancy metadata, and preserve correlation/name semantics.
• csrc/profiler/unavailable_device_tracer.cc: provide the no-device-activity fallback factory.
• tests/unit/test_profiler_privateuse1.py: make profiler library checks MetaX-aware, include MetaX in CUPTI-compatible build gating, and translate modified comments to English.
• docs/architecture/profiler.md: document tracer selection, MetaX MCPTI behavior, measured parity, and known scanner limits.
• docs/reference/compatibility.md: record MetaX profiler support as experimental based on local C550 measurement rather than CI validation.

Changes by commit:

1. b349929 - feat: add complete MetaX profiler tracing: adds the MetaX MCPTI implementation, fallback tracer, tests, and documentation.

Verification

Pre-submission Checklist

• Linting passed (ruff check, ruff format --check)
• Type checking passed (not applicable; C++ and pytest changes)
• All tests pass (relevant unit and integration tests; one pre-existing conditional unit test remains skipped)
• Manual testing completed (MetaX C550 profiler workload reproduced and verified)
• No debug/temporary code (diagnostics are environment-gated and no dump instrumentation remains)
• Documentation updated (profiler architecture and compatibility reference)
• Commit messages follow conventions (type: description format)
• All text in English

Linting Results

$ ruff --version
ruff 0.15.12
$ ruff check .
All checks passed!
$ ruff format --check .
176 files already formatted

Test Results

# Command:
env ACCELERATOR=metax METAX_PATH=/opt/maca-3.8.0 FLAGOS_METAX_BOXING=1 \
  LD_LIBRARY_PATH=/opt/maca-3.8.0/lib:$LD_LIBRARY_PATH \
  PYTHONPATH=$PWD \
  /public-flash/lvyufeng/miniconda3/envs/maca-torch210-py312/bin/python \
  -m pytest tests/integration/test_profiler_parity.py -v -m main_ops

# Output:
7 passed, 1 warning in 0.71s

# Command:
env ACCELERATOR=metax METAX_PATH=/opt/maca-3.8.0 FLAGOS_METAX_BOXING=1 \
  LD_LIBRARY_PATH=/opt/maca-3.8.0/lib:$LD_LIBRARY_PATH \
  PYTHONPATH=$PWD \
  /public-flash/lvyufeng/miniconda3/envs/maca-torch210-py312/bin/python \
  -m pytest tests/unit/test_profiler_privateuse1.py -v

# Output:
4 passed, 1 skipped, 2 warnings in 11.57s

Manual Verification

# Command:
env FLAGOS_METAX_BOXING=1 cmake --build build --target torch_fl -j 8

# Output:
ninja: no work to do.

# Command:
/usr/bin/c++ -std=c++17 -fPIC \
  -I$PWD/csrc/profiler -I$PWD/csrc \
  -I/public-flash/lvyufeng/miniconda3/envs/maca-torch210-py312/lib/python3.12/site-packages/torch/include \
  -c csrc/profiler/unavailable_device_tracer.cc \
  -o /tmp/unavailable_device_tracer.o

# Output:
completed successfully with exit code 0

The original MetaX parity workload initially produced no usable device events because MCPTI records were misnamed, malformed iterator advancement abandoned valid records, and zero/zero timestamps were converted into an unrelated clock range. After the implementation, the same workload emitted kernel=17, memcpy=1, memset=1, and runtime=187 events; all seven parity assertions passed.

Code Quality Verification

Style Consistency

• Matched existing code style in modified files
• Followed naming conventions and existing tracer interfaces
• Kept comments focused on vendor-specific behavior and failure modes
• Reused the existing Kineto and DeviceTracer abstractions

Edge Cases Considered

1. MetaX and NVIDIA callback-ID namespaces are different.
2. MCPTI may return a repeated or stale iterator pointer.
3. Vendor records may have different complete sizes than the mirrored common prefixes.
4. start == end == 0 means timing was unavailable, not a valid event at clock origin.
5. Torch external correlation ID zero is valid and must remain distinguishable from absent.
6. Platforms without CUPTI, MCPTI, ROCtracer, or MSPTI must still link through the unavailable tracer.

Potential Risks

1. The MetaX scanner is measured on MACA 3.8.0 and has not been validated across all MetaX SDK versions, devices, or non-default streams.
2. DCU and BPU alternate hardware paths could not be rebuilt on this host because their vendor SDKs are not installed.
3. CUPTI-compatible activity record mirrors remain version-sensitive and are intentionally guarded by layout checks.

Rollback Plan

Revert commit b349929. This removes the MetaX MCPTI parser changes, fallback tracer, source-selection changes, tests, and documentation updates without changing generated operator code.

Related Work

• Related to the existing CUDA profiler parity implementation.
• Related to the existing DCU ROCtracer implementation and upstream Ascend MSPTI source selection.
• No issue number was provided for this change.

Explicitly Not Included

• No new handwritten vendor kernels.
• No changes to distributed communication or operator routing.
• No claims of CI validation for MetaX profiler support.
• No DCU, BPU, or PPU hardware validation beyond source-selection coverage available on this host.

Human Review Notes

Areas needing special attention:

1. Verify the MetaX MCPTI record scanner against additional MACA SDK versions and non-default streams.
2. Review the deferred mcptiActivityGetApiName resolution and callback ABI typedefs for SDK compatibility.
3. Review the CMake tracer-selection matrix alongside the Ascend MSPTI and DCU ROCtracer implementations.

Questions for reviewer:

1. Should a MetaX-specific scanner fixture be added once a second MCPTI SDK/device environment is available?
2. Should the measured C550 parity command become a vendor CI job when a MetaX runner is available?

Additional Context

• The branch was rebased onto flagos/main before submission.
• The upstream-compatible target repository is flagos-ai/Torch-FL.

---

🤖 Generated with Claude Code