feat(pipeline_parallel): support single-host pipeline parallelism for Llama

tpu_inference/envs.py

@@ -82,6 +82,7 @@
     MOE_ROUTE_PADDING_TO_EXPERT0: bool = False
     VLLM_TPU_BUCKET_PADDING_GAP: int = 0
     VLLM_INCREMENTAL_FP8_LOADING: bool = False
+    VLLM_INCREMENTAL_MXFP4_LOADING: bool = False
     TPU_MESH_SORT_BY_COORDS: bool = False
 
 
@@ -494,6 +495,10 @@ def _get_int_list_env() -> list[int]:
     # when initializing large FP8 models on smaller RAM TPUs such as TPU8i.
     "VLLM_INCREMENTAL_FP8_LOADING":
     env_bool("VLLM_INCREMENTAL_FP8_LOADING", default=False),
+    # Controls whether MXFP4 MoE layers perform incremental weight
+    # loading, sharding, and immediate host RAM cleanup.
+    "VLLM_INCREMENTAL_MXFP4_LOADING":
+    env_bool("VLLM_INCREMENTAL_MXFP4_LOADING", default=False),
 }
 
 

tpu_inference/layers/vllm/quantization/base.py

@@ -13,14 +13,67 @@
 # limitations under the License.
 
 from abc import ABC, abstractmethod
+import ctypes
+import ctypes.util
+import gc
+from typing import Optional
 
+import jax
 import torch
 from vllm.logger import init_logger
 from vllm.model_executor.layers import linear as vllm_linear
 
+from tpu_inference import envs
+
 logger = init_logger(__name__)
 
 
+def _free_torch_storage(tensor: Optional[torch.Tensor]) -> None:
+    """Safely frees the underlying CPU memory storage of a PyTorch tensor.
+
+    Tries `untyped_storage().resize_(0)` first, with fallback to `set_(torch.storage.UntypedStorage())`
+    for 0-dim scalars or float8 dtypes that cannot be resized in-place.
+    """
+    if tensor is None:
+        return
+    try:
+        tensor.untyped_storage().resize_(0)
+    except Exception:
+        try:
+            tensor.set_(torch.storage.UntypedStorage())
+        except Exception:
+            pass
+
+
+def _release_host_memory() -> None:
+    """Frees CPU host memory and trims malloc arena if incremental loading is enabled."""
+    if not (getattr(envs, "VLLM_INCREMENTAL_FP8_LOADING", False) or getattr(envs, "VLLM_INCREMENTAL_MXFP4_LOADING", False)):
+        return
+    gc.collect()
+    jax.effects_barrier()
+    try:
+        libc_name = ctypes.util.find_library("c")
+        if libc_name:
+            ctypes.CDLL(libc_name).malloc_trim(0)
+    except Exception as e:
+        logger.debug(f"malloc_trim failed: {e}")
+
+
+def _log_memory_stats(layer_name: str = "") -> None:
+    try:
+        import psutil, resource
+        proc = psutil.Process()
+        rss_gb = proc.memory_info().rss / (1024 ** 3)
+        max_rss_gb = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / (1024 ** 2)
+        print(
+            f"[RAM Trace] Layer {layer_name} sharded & freed | "
+            f"Process RSS: {rss_gb:.2f} GB | Peak RSS: {max_rss_gb:.2f} GB",
+            flush=True,
+        )
+    except Exception as e:
+        print(f"[RAM Trace Error] {e}", flush=True)
+
+
 class VllmQuantizationMethod(ABC):
 
     def maybe_process_linear_weights(

tpu_inference/layers/vllm/quantization/fp8.py

@@ -51,7 +51,8 @@
     select_moe_backend_from_fused_moe_config, vllm_moe_apply)
 from tpu_inference.layers.vllm.process_weights.cleanup_sharding import \
     _tensor_is_in_cpu
-from tpu_inference.layers.vllm.quantization.base import VllmQuantizationMethod
+from tpu_inference.layers.vllm.quantization.base import (
+    VllmQuantizationMethod, _log_memory_stats)
 from tpu_inference.layers.vllm.quantization.configs import (
     VllmQuantConfig, VllmQuantLinearConfig)
 from tpu_inference.layers.vllm.quantization.unquantized import (
@@ -324,6 +325,7 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
                 layer.bias = to_parameter_list(weights.bias)
 
         _release_host_memory()
+        _log_memory_stats(layer_name=getattr(layer, "_module_name", getattr(layer, "prefix", str(type(layer)))))
 
     def apply(self,
               layer: torch.nn.Module,
@@ -473,26 +475,32 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
 
         del w13_weight, w2_weight, w13_weight_scale, w2_weight_scale, input_weights
 
-        weights = torch_view(
-            shard_moe_weights(weights, self.moe_backend, self.mesh))
+        sharded = shard_moe_weights(weights, self.moe_backend, self.mesh)
+        del weights
+
+        tv_weights = torch_view(sharded)
+        del sharded
 
-        layer.w13_weight = Parameter(weights.w13_weight, requires_grad=False)
-        layer.w2_weight = Parameter(weights.w2_weight, requires_grad=False)
+        layer.w13_weight = Parameter(tv_weights.w13_weight, requires_grad=False)
+        layer.w2_weight = Parameter(tv_weights.w2_weight, requires_grad=False)
 
         # Use setattr to dynamically assign the correct scale parameter name
         # based on the quantization type. vLLM uses 'weight_scale_inv' for
         # block-quantized scales and 'weight_scale' for per-tensor/per-channel scales.
         setattr(
             layer,
             scale_w13_name,
-            Parameter(weights.w13_weight_scale, requires_grad=False),
+            Parameter(tv_weights.w13_weight_scale, requires_grad=False),
         )
         setattr(
             layer,
             scale_w2_name,
-            Parameter(weights.w2_weight_scale, requires_grad=False),
+            Parameter(tv_weights.w2_weight_scale, requires_grad=False),
         )
+        del tv_weights
+
         _release_host_memory()
+        _log_memory_stats(layer_name=getattr(layer, "_module_name", getattr(layer, "prefix", str(type(layer)))))
 
     def apply_monolithic(
         self,