Optimize offload weight pinning with mmap host registration

[Fatemanx]

Summary

This PR optimizes CPU offload weight loading by registering eligible CPU
tensor storage as CUDA pinned host memory in place.

For lazy-loaded safetensors weights, the previous path allocated a pinned CPU
tensor and copied the source CPU tensor into it. This change allows the
original CPU storage to be used directly as the pinned H2D source when no
dtype conversion is required, avoiding one extra host-side copy.

Changes

• Add shared pinned-weight helpers for CUDA host registration and fallback
  handling.
• Prefer in-place cudaHostRegister for eligible CPU tensors.
• Fall back to the existing copy-based pinned allocation path when
  registration is unavailable or dtype conversion is required.
• Mark mmap-backed registered tensors as readonly CPU sources so to_cpu()
  does not write device data back into file-backed storage.
• Route MM, norm, embedding, tensor, and CUDA platform weight paths through
  the shared helper.
• Preserve existing dtype semantics for lazy reload paths, including fp32
  scale/bias cases.
• Add unit tests for registration, fallback, transpose views, dtype
  conversion, lazy reload, and CPU offload copy-back behavior.

Validation

CUDA_VISIBLE_DEVICES=1 PYTHONPATH=. python -m unittest
test_cases.test_mmap_pinned_weights

Ran 11 tests OK

Notes

The in-place registered CPU source is treated as readonly because offloaded
inference weights are immutable. Copy-based pinned buffers keep the previous
writable copy-back behavior.

[gemini-code-assist]

Code Review

This pull request introduces an optimization to register existing CPU tensor storage as CUDA pinned host memory in-place using cudaHostRegister, avoiding unnecessary memory allocations and copies. It refactors embedding, matrix multiplication, normalization, and tensor classes to utilize the new create_pin_tensor and move_tensor_back_to_cpu utility functions, and adds a comprehensive test suite. Feedback was provided to make the unregister finalizer in lightx2v/common/ops/utils.py more robust during interpreter shutdown by capturing torch functions as default arguments to avoid potential AttributeError or NameError exceptions.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.