Support dynamic shard_size when save_only_first_rank = true

[larksudo]

1. Support dynamic shard_size when save_only_first_rank = true

python -m pytest tests/v1/test_estimate_shard_size.py -v
=============================================================== test session starts ================================================================
platform linux -- Python 3.11.14, pytest-9.1.1, pluggy-1.6.0 -- /usr/local/python3.11.14/bin/python
cachedir: .pytest_cache
rootdir: /home/xnm/vl/LMCache-Ascend/tests
configfile: pytest.ini
plugins: anyio-4.13.0
collected 11 items                                                                                                                                 

tests/v1/test_estimate_shard_size.py::TestEstimateShardSize::test_normal_case PASSED                                                         [  9%]
tests/v1/test_estimate_shard_size.py::TestEstimateShardSize::test_moderate_memory PASSED                                                     [ 18%]
tests/v1/test_estimate_shard_size.py::TestEstimateShardSize::test_zero_free_memory PASSED                                                    [ 27%]
tests/v1/test_estimate_shard_size.py::TestEstimateShardSize::test_tiny_free_memory PASSED                                                    [ 36%]
tests/v1/test_estimate_shard_size.py::TestEstimateShardSize::test_huge_chunk_size PASSED                                                     [ 45%]
tests/v1/test_estimate_shard_size.py::TestEstimateShardSize::test_per_chunk_larger_than_available PASSED                                     [ 54%]
tests/v1/test_estimate_shard_size.py::TestEstimateShardSize::test_per_chunk_zero_raises PASSED                                               [ 63%]
tests/v1/test_estimate_shard_size.py::TestEstimateShardSize::test_boundary_values[42949672960-0-1073741824-5] PASSED                         [ 72%]
tests/v1/test_estimate_shard_size.py::TestEstimateShardSize::test_boundary_values[25769803776-0-1073741824-3] PASSED                         [ 81%]
tests/v1/test_estimate_shard_size.py::TestEstimateShardSize::test_boundary_values[34359738368-0-268435456-16] PASSED                         [ 90%]
tests/v1/test_estimate_shard_size.py::TestEstimateShardSize::test_boundary_values[141733920768-4294967296-1073741824-16] PASSED              [100%]

================================================================= warnings summary =================================================================
../../../../usr/local/python3.11.14/lib/python3.11/site-packages/opentelemetry/sdk/_logs/_internal/__init__.py:531
../../../../usr/local/python3.11.14/lib/python3.11/site-packages/opentelemetry/sdk/_logs/_internal/__init__.py:531
../../../../usr/local/python3.11.14/lib/python3.11/site-packages/opentelemetry/sdk/_logs/_internal/__init__.py:531
../../../../usr/local/python3.11.14/lib/python3.11/site-packages/opentelemetry/sdk/_logs/_internal/__init__.py:531
../../../../usr/local/python3.11.14/lib/python3.11/site-packages/opentelemetry/sdk/_logs/_internal/__init__.py:531
../../../../usr/local/python3.11.14/lib/python3.11/site-packages/opentelemetry/sdk/_logs/_internal/__init__.py:531
v1/test_estimate_shard_size.py::TestEstimateShardSize::test_normal_case
  /usr/local/python3.11.14/lib/python3.11/site-packages/opentelemetry/sdk/_logs/_internal/__init__.py:531: DeprecationWarning: `LoggingHandler` in `opentelemetry-sdk` is deprecated. Use the handler from `opentelemetry-instrumentation-logging` instead.
    warnings.warn(

../../../../usr/local/python3.11.14/lib/python3.11/site-packages/torch_npu/contrib/transfer_to_npu.py:347
  /usr/local/python3.11.14/lib/python3.11/site-packages/torch_npu/contrib/transfer_to_npu.py:347: ImportWarning: 
      *************************************************************************************************************
      The torch.Tensor.cuda and torch.nn.Module.cuda are replaced with torch.Tensor.npu and torch.nn.Module.npu now..
      The torch.cuda.DoubleTensor is replaced with torch.npu.FloatTensor cause the double type is not supported now..
      The backend in torch.distributed.init_process_group set to hccl now..
      The torch.cuda.* and torch.cuda.amp.* are replaced with torch.npu.* and torch.npu.amp.* now..
      The device parameters have been replaced with npu in the function below:
      torch.logspace, torch.randint, torch.hann_window, torch.rand, torch.full_like, torch.ones_like, torch.rand_like, torch.randperm, torch.arange, torch.frombuffer, torch.normal, torch._empty_per_channel_affine_quantized, torch.empty_strided, torch.empty_like, torch.scalar_tensor, torch.tril_indices, torch.bartlett_window, torch.ones, torch.sparse_coo_tensor, torch.randn, torch.kaiser_window, torch.tensor, torch.triu_indices, torch.as_tensor, torch.zeros, torch.randint_like, torch.full, torch.eye, torch._sparse_csr_tensor_unsafe, torch.empty, torch._sparse_coo_tensor_unsafe, torch.blackman_window, torch.zeros_like, torch.range, torch.sparse_csr_tensor, torch.randn_like, torch.from_file, torch._cudnn_init_dropout_state, torch._empty_affine_quantized, torch.linspace, torch.hamming_window, torch.empty_quantized, torch._pin_memory, torch.load, torch.set_default_device, torch.get_device_module, torch.sparse_compressed_tensor, torch.Tensor.new_empty, torch.Tensor.new_empty_strided, torch.Tensor.new_full, torch.Tensor.new_ones, torch.Tensor.new_tensor, torch.Tensor.new_zeros, torch.Tensor.to, torch.Tensor.pin_memory, torch.nn.Module.to, torch.nn.Module.to_empty
      *************************************************************************************************************
      
    warnings.warn(msg, ImportWarning)

../../../../usr/local/python3.11.14/lib/python3.11/site-packages/torch_npu/contrib/transfer_to_npu.py:276
  /usr/local/python3.11.14/lib/python3.11/site-packages/torch_npu/contrib/transfer_to_npu.py:276: RuntimeWarning: torch.jit.script and torch.jit.script_method will be disabled by transfer_to_npu, which currently does not support them, if you need to enable them, please do not use transfer_to_npu.
    warnings.warn(msg, RuntimeWarning)

<frozen importlib._bootstrap>:241
  <frozen importlib._bootstrap>:241: DeprecationWarning: builtin type SwigPyPacked has no __module__ attribute

<frozen importlib._bootstrap>:241
  <frozen importlib._bootstrap>:241: DeprecationWarning: builtin type SwigPyObject has no __module__ attribute

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
========================================================= 11 passed, 11 warnings in 2.02s ==========================================================
sys:1: DeprecationWarning: builtin type swigvarlink has no __module__ attribute
root@liteserver-hps-805f-00002:/home/xnm/vl/LMCache-Ascend# 

[gemini-code-assist]

Code Review

This pull request introduces auto-estimation and validation of broadcast_shard_size based on available NPU memory to prevent OOM errors on Ascend devices, along with preallocation of the broadcast merged-buffer pool and explicit stream synchronization to prevent hangs in graph mode. The review feedback highlights several key improvements for robustness and compatibility, including raising a clear error when load_stream is missing, removing the Python 3.10+ strict argument from zip for backward compatibility, preventing potential division-by-zero and type errors, and simplifying the dtype byte-size calculation using PyTorch's built-in element_size().

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.

[matthewygf]

lgtm