fix(disagg): NVFP4 + publish-only rollout fixes for disaggregated tra…

miles/backends/megatron_utils/megatron_to_hf/processors/__init__.py

@@ -2,13 +2,15 @@
 from .quantizer_compressed_tensors import quantize_params_compressed_tensors
 from .quantizer_fp8 import quantize_params_fp8
 from .quantizer_mxfp8 import quantize_params_mxfp8
+from .quantizer_nvfp4 import quantize_params_nvfp4
 
 __all__ = [
     "remove_padding",
     "quantize_param",
     "quantize_params_fp8",
     "quantize_params_mxfp8",
     "quantize_params_compressed_tensors",
+    "quantize_params_nvfp4",
 ]
 
 
@@ -22,3 +24,7 @@ def quantize_params(args, megatron_name, converted_named_params, quantization_co
     elif quantization_config["quant_method"] == "compressed-tensors":
         # only int4 at the moment.
         return quantize_params_compressed_tensors(converted_named_params, quantization_config)
+    elif quantization_config.get("quant_algo") == "NVFP4" or quantization_config.get("quant_method") == "nvfp4":
+        # modelopt NVFP4 checkpoints (e.g. nvidia/Kimi-K2.6-NVFP4) advertise
+        # quant_method="modelopt", so route on quant_algo, not quant_method.
+        return quantize_params_nvfp4(args, megatron_name, converted_named_params, quantization_config)

miles/backends/megatron_utils/update_weight/update_weight_from_disk_delta.py

@@ -143,7 +143,8 @@ def _snapshot_bucket(self, converted_named_tensors: list[tuple[str, torch.Tensor
             try:
                 self._snapshot[name] = self._baseline_reader(name)
             except KeyError:
-                self._snapshot[name] = tensor.detach().cpu().contiguous().view(torch.uint8).numpy().reshape(-1)
+                # flatten before the byte-view (0-dim scalars like NVFP4 scales)
+                self._snapshot[name] = tensor.detach().cpu().contiguous().reshape(-1).view(torch.uint8).numpy()
                 logger.warning("seed: %s absent from hf_checkpoint; seeding from current weights", name)
         converted_named_tensors.clear()
 
@@ -196,7 +197,11 @@ def _encode_bucket(self, converted_named_tensors: list[tuple[str, torch.Tensor]]
         """Copy each gathered HF tensor to host and submit it to the diff/compress pool, draining
         once enough work is in flight to backpressure the gather (source ranks only)."""
         for name, tensor in converted_named_tensors:
-            flat = tensor.detach().contiguous().view(torch.uint8).reshape(-1)
+            # Flatten to 1-D *before* the byte-view: NVFP4 export emits 0-dim
+            # scalar tensors (weight_scale_2, input_scale), and .view(torch.uint8)
+            # cannot reinterpret a 0-dim tensor ("self.dim() cannot be 0 to view
+            # Float as Byte"). reshape(-1) on a 0-dim tensor yields [1] first.
+            flat = tensor.detach().contiguous().reshape(-1).view(torch.uint8)
             nbytes = int(flat.numel())
             if self._use_pinned and nbytes <= self._max_bytes:
                 buf = self._free_q.get()  # blocks when all buffers are in flight -> backpressures the gather

miles/rollout/sglang_rollout.py

@@ -107,9 +107,17 @@ def __init__(self, args: Namespace) -> None:
         )
         self.processor = load_processor(args.hf_checkpoint, trust_remote_code=True)
 
-        self.semaphore = asyncio.Semaphore(
-            args.sglang_server_concurrency * args.rollout_num_gpus // args.rollout_num_gpus_per_engine
-        )
+        if getattr(args, "rollout_endpoint_url", None):
+            # Publish-only / disaggregated: no miles-launched engines, so
+            # rollout_num_gpus == 0 and the GPU-based formula below would be 0,
+            # making asyncio.Semaphore(0) deadlock every rollout. Bound generation
+            # concurrency by the external pool's per-engine concurrency instead.
+            generation_concurrency = args.sglang_server_concurrency
+        else:
+            generation_concurrency = (
+                args.sglang_server_concurrency * args.rollout_num_gpus // args.rollout_num_gpus_per_engine
+            )
+        self.semaphore = asyncio.Semaphore(generation_concurrency)
         self.sampling_params: dict[str, Any] = dict(
             temperature=args.rollout_temperature,
             top_p=args.rollout_top_p,

miles/utils/chat_template_utils/deepseek_v4.py

@@ -7,7 +7,10 @@
 import os
 from typing import Any
 
-from sglang.srt.entrypoints.openai import encoding_dsv4
+try:
+    from sglang.srt.entrypoints.openai import encoding_dsv4
+except ImportError:  # older sglang-miles without the DeepSeek-V4 encoder; fine for non-V4 models
+    encoding_dsv4 = None
 from sglang.srt.entrypoints.openai.protocol import Tool
 
 logger = logging.getLogger(__name__)

miles/utils/http_utils.py

@@ -227,10 +227,18 @@ async def _post(client, url, payload, max_retries=60, action="post", headers=Non
 def init_http_client(args):
     """Initialize HTTP client and optionally enable distributed POST via Ray."""
     global _http_client, _client_concurrency, _distributed_post_enabled
-    if not args.rollout_num_gpus:
+    publish_only = bool(getattr(args, "rollout_endpoint_url", None))
+    if not args.rollout_num_gpus and not publish_only:
         return
 
-    _client_concurrency = args.sglang_server_concurrency * args.rollout_num_gpus // args.rollout_num_gpus_per_engine
+    if publish_only:
+        # Disaggregated / publish-only: no miles-launched engines, so
+        # rollout_num_gpus == 0 and the GPU-based formula below would be 0 — the
+        # client would never be created (NoneType.post). The external pool bounds
+        # concurrency instead.
+        _client_concurrency = args.sglang_server_concurrency
+    else:
+        _client_concurrency = args.sglang_server_concurrency * args.rollout_num_gpus // args.rollout_num_gpus_per_engine
     if _http_client is None:
         _http_client = httpx.AsyncClient(
             limits=httpx.Limits(max_connections=_client_concurrency),