feat(inference): opt-in online FP8 ue8m0 scales and stochastic weight rounding

packages/prime-rl-configs/src/prime_rl/configs/inference.py

@@ -31,6 +31,15 @@ class ServerConfig(BaseConfig):
 
 QuantizationType = Literal["fp8_per_block"]
 
+_ONLINE_PER_BLOCK_FP8_NAMES = frozenset(("fp8_per_block", "fp8_per_block_static"))
+
+
+def _selects_online_per_block_fp8(spec: Any) -> bool:
+    """Whether a vLLM ``quantization_config`` spec (a name or a ``{weight, activation}`` dict) quantizes weights to per-block FP8."""
+    if isinstance(spec, dict):
+        spec = spec.get("weight")
+    return spec in _ONLINE_PER_BLOCK_FP8_NAMES
+
 
 class VllmConfig(BaseConfig):
     """Arguments forwarded to the vLLM server, under vLLM's own argument names
@@ -113,6 +122,20 @@ class VllmConfig(BaseConfig):
     enable_return_routed_experts: bool = False
     """Return routed experts in responses."""
 
+    def uses_online_per_block_fp8(self) -> bool:
+        """Whether vLLM quantizes linear or MoE weights online to per-block FP8, resolved as in vLLM's ``resolve_quantization_config``."""
+        quantization_config = (self.model_extra or {}).get("quantization_config") or {}
+        targets = quantization_config.get("targets") or {}
+        if any(_selects_online_per_block_fp8(spec) for spec in targets.values()):
+            return True
+        for layer_kind in ("linear", "moe"):
+            spec = quantization_config.get(layer_kind)
+            if spec is None:
+                spec = self.quantization
+            if _selects_online_per_block_fp8(spec):
+                return True
+        return False
+
     @model_validator(mode="before")
     @classmethod
     def parse_extra_values(cls, data: dict) -> dict:
@@ -430,6 +453,12 @@ class InferenceConfig(BaseConfig):
     use_deep_gemm: bool = False
     """Enable vLLM DeepGEMM FP8 kernels ``VLLM_USE_DEEP_GEMM=1``. Only works with block-wise FP8 quantization (e.g. GLM-5-FP8)."""
 
+    fp8_ue8m0_weight_scales: bool = False
+    """Quantize online block-wise FP8 weights with power-of-two (UE8M0) block scales instead of ``amax / 448``. Weights of a dequantized FP8 release (e.g. ``PrimeIntellect/DeepSeek-V4-Flash-0731-bf16``) already sit on the e4m3 grid under power-of-two scales, so they round-trip exactly instead of picking up e4m3 rounding error. Sets ``PRIME_FP8_UE8M0_WEIGHT_SCALES=1`` for the vLLM workers."""
+
+    fp8_stochastic_weight_rounding: bool = False
+    """Round online block-wise FP8 weights stochastically instead of to nearest. Each element lands on the e4m3 grid point on the far side of its value with probability equal to its fractional distance in bins, so the served weight equals the bf16 weight in expectation. Round-to-nearest pins a weight that starts on the grid until the trainer has moved it half a bin, which at lr 1e-6 takes on the order of 1000 steps, so the served policy stays at step 0 while the trainer drifts. Unbiased rounding lets it track sub-bin updates in expectation. Input already on the grid rounds to nearest, so the checkpoint quantizes to identical bytes. Meant for serving a policy under RL weight updates: for evaluating a fixed checkpoint it only serves one noisy draw of the weights. Sets ``PRIME_FP8_STOCHASTIC_WEIGHT_ROUNDING=1`` for the vLLM workers."""
+
     weight_broadcast: WeightBroadcastConfig = WeightBroadcastConfig()
 
     kv_cache_offload: KVCacheOffloadConfig | None = None
@@ -466,6 +495,20 @@ def validate_multi_node_requires_slurm(self):
             raise ValueError("Must use SLURM for multi-node / disaggregated deployment.")
         return self
 
+    @model_validator(mode="after")
+    def validate_online_fp8_rounding_requires_online_per_block_fp8(self):
+        """The FP8 weight-rounding options patch vLLM's online per-block FP8 caster, so they do nothing without it."""
+        enabled = [
+            name for name in ("fp8_ue8m0_weight_scales", "fp8_stochastic_weight_rounding") if getattr(self, name)
+        ]
+        if enabled and not self.vllm.uses_online_per_block_fp8():
+            raise ValueError(
+                f"{' and '.join(enabled)} can only take effect with online per-block FP8 weight quantization. Set "
+                'inference.vllm.quantization = "fp8_per_block" or select fp8_per_block for linear or MoE layers '
+                "in inference.vllm.quantization_config."
+            )
+        return self
+
     @model_validator(mode="after")
     def validate_llmd_no_routed_experts(self):
         """Reject routed-expert return with the llm-d router (breaks P/D, unverified for multi-node)."""

packages/prime-rl-configs/src/prime_rl/configs/rl.py

@@ -21,6 +21,8 @@
     WeightBroadcastConfig,
 )
 from prime_rl.configs.trainer import (
+    DeepGemmFP8MoEComputeConfig,
+    FP8Config,
     TokenizerConfig,
     TrainerConfig,
 )
@@ -474,6 +476,21 @@ def validate_eplb(self):
             raise ValueError("inference.vllm.enable_eplb is not supported with RL weight updates.")
         return self
 
+    @model_validator(mode="after")
+    def validate_stochastic_fp8_rounding_with_fp8_trainer(self):
+        """Reject inference-side stochastic FP8 weight rounding with an FP8 trainer, without checking which layers each side quantizes."""
+        if self.inference is None or not self.inference.fp8_stochastic_weight_rounding:
+            return self
+        model = self.trainer.model
+        if isinstance(model.quantization, FP8Config) or isinstance(model.moe.compute, DeepGemmFP8MoEComputeConfig):
+            raise ValueError(
+                "inference.fp8_stochastic_weight_rounding is not supported with an FP8 trainer "
+                "(trainer.model.quantization of type fp8 or trainer.model.moe.compute of type deepgemm_fp8). "
+                "The trainer re-quantizes its weights with round-to-nearest every step, so stochastic rounding "
+                "on the inference side only adds trainer/inference mismatch."
+            )
+        return self
+
     @model_validator(mode="after")
     def auto_setup_lora(self):
         if self.trainer.model.lora is not None:

src/prime_rl/inference/patches.py

@@ -20,6 +20,8 @@ def apply_shared_vllm_patches():
     monkey_patch_minimax_m2_think_end_passthrough()
     monkey_patch_kv_xfer_finished_tolerate_freed()
     monkey_patch_online_fp8_parameter_cast()
+    monkey_patch_fp8_ue8m0_weight_scales()
+    monkey_patch_fp8_stochastic_weight_rounding()
     monkey_patch_deepseek_v4_allowed_layer_types()
     monkey_patch_tokenize_params_validation()
     monkey_patch_strip_routed_experts_from_chat()
@@ -154,6 +156,114 @@ def _per_block_cast_to_fp8(x, *args, **kwargs):
     fp8.per_block_cast_to_fp8 = _per_block_cast_to_fp8
 
 
+def monkey_patch_fp8_ue8m0_weight_scales():
+    """Off unless ``PRIME_FP8_UE8M0_WEIGHT_SCALES=1`` (``inference.fp8_ue8m0_weight_scales``): power-of-two block scales.
+
+    vLLM's online block-FP8 linear and MoE methods call
+    ``per_block_cast_to_fp8(..., use_ue8m0=False)``, which picks the scale ``amax / 448``. A
+    dequantized FP8 release already lies on the e4m3 grid under power-of-two block scales, so
+    forcing ``use_ue8m0=True`` restores that grid and the weights round-trip exactly. The scales
+    stay fp32, so no kernel change is involved.
+    """
+    if os.environ.get("PRIME_FP8_UE8M0_WEIGHT_SCALES") != "1":
+        return
+
+    from vllm.logger import init_logger
+    from vllm.model_executor.layers.quantization.online import fp8
+
+    original_cast = fp8.per_block_cast_to_fp8
+    if getattr(original_cast, "_prime_rl_forces_ue8m0", False):
+        return
+
+    def _per_block_cast_to_fp8(x, *args, **kwargs):
+        kwargs["use_ue8m0"] = True
+        return original_cast(x, *args, **kwargs)
+
+    _per_block_cast_to_fp8._prime_rl_forces_ue8m0 = True
+    fp8.per_block_cast_to_fp8 = _per_block_cast_to_fp8
+    init_logger("vllm.prime_rl.fp8").info(
+        "PRIME_FP8_UE8M0_WEIGHT_SCALES=1: quantizing online FP8 weights with power-of-two block scales."
+    )
+
+
+_STOCHASTIC_ROUNDING_ROW_CHUNK = 2048
+_E4M3_SIGN_BIT = 0x80
+_E4M3_MAGNITUDE_MASK = 0x7F
+_E4M3_MAX_FINITE_MAGNITUDE = 0x7E
+
+
+def _expand_block_scales(scales: torch.Tensor, block_size: list[int], rows: int, cols: int) -> torch.Tensor:
+    block_m, block_n = block_size
+    return scales.repeat_interleave(block_m, dim=0)[:rows].repeat_interleave(block_n, dim=1)[:, :cols]
+
+
+def stochastic_round_fp8(x_scaled: torch.Tensor, q_near: torch.Tensor) -> torch.Tensor:
+    """Move each round-to-nearest e4m3 value ``q_near`` to its neighbour on the far side of ``x_scaled`` with probability equal to the fractional distance, so the result equals ``x_scaled`` in expectation."""
+    q_near_float = q_near.float()
+    bits = q_near.contiguous().view(torch.uint8).to(torch.int16)
+    sign = bits & _E4M3_SIGN_BIT
+    magnitude = bits & _E4M3_MAGNITUDE_MASK
+    farther_from_zero = x_scaled.abs() > q_near_float.abs()
+    neighbour_magnitude = torch.where(farther_from_zero, magnitude + 1, magnitude - 1).clamp(
+        0, _E4M3_MAX_FINITE_MAGNITUDE
+    )
+    neighbour = (sign | neighbour_magnitude).to(torch.uint8).view(torch.float8_e4m3fn).float()
+    gap = neighbour - q_near_float
+    probability = torch.where(gap != 0, (x_scaled - q_near_float) / gap, torch.zeros_like(gap)).clamp(0, 1)
+    pick_neighbour = torch.rand_like(probability) < probability
+    return torch.where(pick_neighbour, neighbour, q_near_float).to(torch.float8_e4m3fn)
+
+
+def monkey_patch_fp8_stochastic_weight_rounding():
+    """Off unless ``PRIME_FP8_STOCHASTIC_WEIGHT_ROUNDING=1`` (``inference.fp8_stochastic_weight_rounding``): unbiased weight rounding.
+
+    ``per_block_cast_to_fp8`` rounds to nearest, so a weight that starts on the e4m3
+    grid keeps its served value until the trainer has moved it half a bin, about
+    1000 steps at lr 1e-6. Until then the served policy is pinned at step 0 while
+    the trainer drifts and the trainer-vs-inference mismatch grows. Rounding each
+    weight to the far neighbour with probability equal to its fractional distance
+    makes the served weight unbiased, so it tracks sub-bin updates in expectation.
+    On-grid input has zero fractional distance and rounds to nearest, so the
+    checkpoint itself quantizes to identical bytes.
+
+    Wraps whatever ``per_block_cast_to_fp8`` is installed when this runs, so it
+    composes with the UE8M0 wrapper registered just before it. The block scales
+    come back unchanged. Only the e4m3 payload is re-rounded, in row chunks so
+    the fp32 temporaries stay bounded on large expert weights.
+    """
+    if os.environ.get("PRIME_FP8_STOCHASTIC_WEIGHT_ROUNDING") != "1":
+        return
+
+    from vllm.logger import init_logger
+    from vllm.model_executor.layers.quantization.online import fp8
+    from vllm.utils.deep_gemm import DEFAULT_BLOCK_SIZE
+
+    original_cast = fp8.per_block_cast_to_fp8
+    if getattr(original_cast, "_prime_rl_rounds_stochastically", False):
+        return
+
+    def _per_block_cast_to_fp8(x, *args, **kwargs):
+        quantized, scales = original_cast(x, *args, **kwargs)
+        block_size = kwargs.get("block_size", args[0] if args else DEFAULT_BLOCK_SIZE)
+        block_m = block_size[0]
+        rows, cols = quantized.shape
+        row_chunk = max(block_m, _STOCHASTIC_ROUNDING_ROW_CHUNK // block_m * block_m)
+        out = torch.empty_like(quantized)
+        for row_start in range(0, rows, row_chunk):
+            row_end = min(row_start + row_chunk, rows)
+            scale_rows = scales[row_start // block_m : -(-row_end // block_m)]
+            scale_per_element = _expand_block_scales(scale_rows, block_size, row_end - row_start, cols)
+            x_scaled = x[row_start:row_end].float() * (1.0 / scale_per_element)
+            out[row_start:row_end] = stochastic_round_fp8(x_scaled, quantized[row_start:row_end])
+        return out, scales
+
+    _per_block_cast_to_fp8._prime_rl_rounds_stochastically = True
+    fp8.per_block_cast_to_fp8 = _per_block_cast_to_fp8
+    init_logger("vllm.prime_rl.fp8").info(
+        "PRIME_FP8_STOCHASTIC_WEIGHT_ROUNDING=1: rounding online FP8 weights stochastically."
+    )
+
+
 def monkey_patch_kv_xfer_finished_tolerate_freed():
     """Tolerate KV-transfer finish notifications for already-freed requests.
 

src/prime_rl/inference/server.py

@@ -24,6 +24,8 @@ def setup_vllm_env(config: InferenceConfig):
     deep_gemm_enabled = "1" if config.use_deep_gemm else "0"
     os.environ["VLLM_USE_DEEP_GEMM"] = deep_gemm_enabled
     os.environ["VLLM_MOE_USE_DEEP_GEMM"] = deep_gemm_enabled
+    os.environ["PRIME_FP8_UE8M0_WEIGHT_SCALES"] = "1" if config.fp8_ue8m0_weight_scales else "0"
+    os.environ["PRIME_FP8_STOCHASTIC_WEIGHT_ROUNDING"] = "1" if config.fp8_stochastic_weight_rounding else "0"
 
     if config.vllm.enable_lora:
         os.environ["VLLM_ALLOW_RUNTIME_LORA_UPDATING"] = "True"

tests/unit/inference/test_fp8_stochastic_rounding.py

@@ -0,0 +1,22 @@
+import torch
+
+from prime_rl.inference.patches import stochastic_round_fp8
+
+
+def test_stochastic_round_fp8_preserves_on_grid_values():
+    values = torch.tensor([-2.0, -1.125, -0.5, 0.0, 0.5, 1.125, 2.0])
+    quantized = values.to(torch.float8_e4m3fn)
+
+    rounded = stochastic_round_fp8(values, quantized)
+
+    assert torch.equal(rounded.view(torch.uint8), quantized.view(torch.uint8))
+
+
+def test_stochastic_round_fp8_is_unbiased_at_midpoint():
+    torch.manual_seed(0)
+    midpoint = torch.full((200_000,), 1.0625)
+    nearest = midpoint.to(torch.float8_e4m3fn)
+
+    rounded = stochastic_round_fp8(midpoint, nearest).float()
+
+    assert abs(rounded.mean().item() - midpoint[0].item()) < 1e-3

tests/unit/test_configs.py

@@ -1178,6 +1178,75 @@ def test_combined_replay_uses_v2_runner(monkeypatch):
     assert os.environ["VLLM_USE_V2_MODEL_RUNNER"] == "1"
 
 
+def test_online_fp8_rounding_defaults_off():
+    config = InferenceConfig()
+
+    assert config.fp8_ue8m0_weight_scales is False
+    assert config.fp8_stochastic_weight_rounding is False
+
+
+@pytest.mark.parametrize("option", ["fp8_ue8m0_weight_scales", "fp8_stochastic_weight_rounding"])
+@pytest.mark.parametrize(
+    "vllm",
+    [
+        {},
+        {"quantization_config": {"linear": "fp8_per_tensor", "moe": "fp8_per_channel"}},
+        {
+            "quantization": "fp8_per_block",
+            "quantization_config": {"linear": "fp8_per_tensor", "moe": {"weight": "mxfp8"}},
+        },
+    ],
+)
+def test_online_fp8_rounding_requires_online_per_block_fp8(option, vllm):
+    with pytest.raises(ValidationError, match=f"{option} can only take effect with online per-block FP8"):
+        InferenceConfig(**{option: True, "vllm": vllm})
+
+
+@pytest.mark.parametrize(
+    "vllm",
+    [
+        {"quantization": "fp8_per_block"},
+        {"quantization": "fp8_per_block", "quantization_config": {"linear": "fp8_per_tensor"}},
+        {"quantization_config": {"moe": "fp8_per_block"}},
+        {"quantization_config": {"linear": {"weight": "fp8_per_block_static"}}},
+        {"quantization_config": {"targets": {"re:.*experts.*": "fp8_per_block"}}},
+    ],
+)
+def test_online_fp8_rounding_accepts_online_per_block_fp8(vllm):
+    config = InferenceConfig(fp8_ue8m0_weight_scales=True, fp8_stochastic_weight_rounding=True, vllm=vllm)
+
+    assert config.fp8_ue8m0_weight_scales is True
+    assert config.fp8_stochastic_weight_rounding is True
+
+
+def _rl_config_with_stochastic_fp8_rounding(trainer_model: dict) -> dict:
+    return {
+        "trainer": {"model": trainer_model},
+        "orchestrator": {},
+        "inference": {"fp8_stochastic_weight_rounding": True, "vllm": {"quantization": "fp8_per_block"}},
+    }
+
+
+@pytest.mark.parametrize(
+    "trainer_model",
+    [{"quantization": {"type": "fp8"}}, {"moe": {"compute": {"type": "deepgemm_fp8"}}}],
+)
+def test_stochastic_fp8_rounding_rejects_fp8_trainer(trainer_model):
+    with pytest.raises(ValidationError, match="re-quantizes its weights with round-to-nearest"):
+        RLConfig.model_validate(_rl_config_with_stochastic_fp8_rounding(trainer_model))
+
+
+@pytest.mark.parametrize(
+    "trainer_model",
+    [{}, {"quantization": {"type": "mxfp8"}}, {"moe": {"compute": {"type": "mxfp8"}}}],
+)
+def test_stochastic_fp8_rounding_accepts_non_fp8_trainer(trainer_model):
+    config = RLConfig.model_validate(_rl_config_with_stochastic_fp8_rounding(trainer_model))
+
+    assert config.inference is not None
+    assert config.inference.fp8_stochastic_weight_rounding is True
+
+
 CUSTOM_RENDERER_SOURCE = """
 from typing import Literal