feat(refit): expose mutable MXFP8 destinations

megatron/core/inference/moe/activations.py

@@ -318,4 +318,6 @@ def squared_relu_and_quantize_mxfp8(
         NUM_BLOCKS=NUM_BLOCKS,
     )
 
-    return MXFP8Tensor(data=out_fp8, scale=out_scale.view(torch.float8_e8m0fnu), backend="triton")
+    return MXFP8Tensor(
+        data=out_fp8, scale=out_scale.view(torch.float8_e8m0fnu), dtype=x.dtype, backend="triton"
+    )

megatron/core/inference/moe/permute.py

@@ -728,6 +728,9 @@ def permute_and_quantize_mxfp8(
     )
 
     permuted_mxfp8 = MXFP8Tensor(
-        data=out_fp8, scale=out_scale.view(torch.float8_e8m0fnu), backend="triton"
+        data=out_fp8,
+        scale=out_scale.view(torch.float8_e8m0fnu),
+        dtype=hidden_states.dtype,
+        backend="triton",
     )
     return permuted_mxfp8, permuted_probs, permutation_map, inclusive_expert_offsets

megatron/core/inference/quantization/mxfp8_tensor.py

@@ -1,6 +1,6 @@
 # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Literal, Optional
 
 import torch
@@ -122,6 +122,22 @@ class MXFP8Tensor:
     data: torch.Tensor  # [M, K] fp8_e4m3fn
     scale: torch.Tensor  # 1D swizzled or [M, K // 32] unswizzled scales
     backend: Optional[MXFP8Backend] = None  # quantization and GEMM backend
+    # Keyword-only preserves the pre-existing third positional ``backend`` argument.
+    # This metadata supports tensor-like conversion consumers such as Megatron
+    # Bridge; it does not constrain the precision accepted by ``quantize_``/``copy_``.
+    # Legacy direct constructors have unknown logical dtype until their next
+    # successful update.
+    dtype: Optional[torch.dtype] = field(default=None, kw_only=True)
+
+    @property
+    def shape(self) -> torch.Size:
+        """Shape of the quantized data storage."""
+        return self.data.shape
+
+    @property
+    def device(self) -> torch.device:
+        """Device holding the quantized tensor."""
+        return self.data.device
 
     def size(self, idx: Optional[int] = None):
         """Wrapper for calling self.data.size()"""
@@ -141,12 +157,60 @@ def scale_2d(self, K: Optional[int] = None) -> torch.Tensor:
         padded_cols = n_col_blocks * MXFP8_SCALE_COL_BLOCK
         return self.scale.reshape(-1, padded_cols)
 
+    def quantize_(self, value: torch.Tensor) -> "MXFP8Tensor":
+        """Quantize a logical tensor into existing storage without changing pointers.
+
+        The source dtype is preserved when supported by the configured backend so
+        the quantizer does not introduce an avoidable intermediate downcast.
+        FlashInfer FP32 inputs are converted to BF16 because that backend accepts
+        FP16/BF16 inputs only. For legacy instances with unknown logical dtype, the
+        first successful update records the dtype actually passed to the quantizer.
+        Shape broadcasting is unsupported because MXFP8 scale storage has fixed
+        geometry.
+        """
+        if self.data.ndim != 2:
+            raise ValueError(
+                "In-place MXFP8 updates require 2D destination storage; "
+                f"got shape {tuple(self.data.shape)}."
+            )
+        if value.shape != self.shape:
+            raise ValueError(
+                f"MXFP8 shape mismatch: expected {tuple(self.shape)}, got {tuple(value.shape)}."
+            )
+        if self.backend is None:
+            raise ValueError("Cannot update an MXFP8Tensor without a quantization backend.")
+        value = value.to(device=self.device)
+        if self.backend == "flashinfer" and value.dtype == torch.float32:
+            value = value.to(dtype=torch.bfloat16)
+        value = value.contiguous()
+        quantized = MXFP8Tensor.from_bf16(value, backend=self.backend)
+        self.data.copy_(quantized.data)
+        self.scale.view(torch.uint8).copy_(quantized.scale.view(torch.uint8))
+        if self.dtype is None:
+            self.dtype = value.dtype
+        return self
+
+    def copy_(self, value: torch.Tensor) -> "MXFP8Tensor":
+        """Tensor-compatible alias for :meth:`quantize_` used by generic writeback.
+
+        Unlike ``torch.Tensor.copy_``, this method requires an exact shape and
+        does not accept ``non_blocking``; quantization and storage updates are
+        ordered on the current CUDA stream. The alias lets external conversion
+        integrations such as Megatron Bridge treat plain and MXFP8 destinations
+        uniformly while ``quantize_`` remains available to existing callers.
+        """
+        return self.quantize_(value)
+
     @classmethod
     def from_bf16(cls, x: torch.Tensor, group_size: int = 32, backend: MXFP8Backend = "flashinfer"):
-        """Quantize BF16 tensor to MXFP8.
+        """Quantize a floating-point CUDA tensor to MXFP8.
+
+        The historical method name is retained for compatibility. The Triton
+        backend accepts BF16, FP16, and FP32 inputs; FlashInfer accepts BF16 and
+        FP16 inputs.
 
         Args:
-            x: [M, K] BF16 tensor on CUDA.
+            x: [M, K] floating-point tensor on CUDA.
             group_size: MXFP8 group size (default 32).
             backend: 'triton' (fused quantize + swizzle Triton kernel) or
                      'flashinfer' (single fused FlashInfer CUDA kernel).
@@ -168,4 +232,4 @@ def from_bf16(cls, x: torch.Tensor, group_size: int = 32, backend: MXFP8Backend
                 f"Unknown MXFP8 quantization backend: '{backend}'. "
                 "Must be 'triton' or 'flashinfer'."
             )
-        return cls(data=data, scale=scale, backend=backend)
+        return cls(data=data, scale=scale, backend=backend, dtype=x.dtype)

megatron/core/inference/quantization/utils.py

@@ -241,9 +241,7 @@ def quantize_params_to_mxfp8(
                     expected_backend=backend,
                     tensor_name=f"persistent MXFP8 parameter {fqn!r}",
                 )
-                new_tensor = MXFP8Tensor.from_bf16(bf16_data, backend=backend)
-                persistent_tensor.data.copy_(new_tensor.data)
-                persistent_tensor.scale.view(torch.uint8).copy_(new_tensor.scale.view(torch.uint8))
+                persistent_tensor.copy_(bf16_data)
                 mcore_tensor = persistent_tensor
             else:
                 # First call: create new MXFP8Tensor

megatron/core/resharding/__init__.py

@@ -1,5 +1,5 @@
 # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
-from .execution import execute_reshard_plan
+from .execution import execute_reshard_plan, refresh_module_caches
 from .planner import (
     build_centralized_reshard_plan,
     build_local_reshard_plan,
@@ -21,6 +21,7 @@
     "build_plan_from_rosters",
     "index_metadata_rosters",
     "execute_reshard_plan",
+    "refresh_module_caches",
     "MXFP8ReshardTransform",
     "ReshardTransform",
     "swap_model_weights",

megatron/core/resharding/execution.py

@@ -18,13 +18,23 @@
 logger = logging.getLogger(__name__)
 
 
-def _refresh_module_caches(dst_module: Optional[torch.nn.Module]) -> None:
-    """Refresh parameter-derived caches in the destination module."""
+def refresh_module_caches(
+    dst_module: torch.nn.Module | list[torch.nn.Module] | tuple[torch.nn.Module, ...] | None,
+) -> None:
+    """Refresh parameter-derived caches in the destination module(s).
+
+    Lists and tuples let external refit callers, including Megatron Bridge,
+    pass virtual-pipeline model chunks directly. ``None`` is a no-op for
+    send-only ranks. The ``dst_module`` parameter name remains stable for
+    callers that use keyword arguments.
+    """
     if dst_module is None:
         return
-    for module in dst_module.modules():
-        if isinstance(module, MegatronModule):
-            module.refresh_cache()
+    roots = dst_module if isinstance(dst_module, (list, tuple)) else (dst_module,)
+    for root in roots:
+        for module in root.modules():
+            if isinstance(module, MegatronModule):
+                module.refresh_cache()
 
 
 @dataclass
@@ -103,7 +113,7 @@ def execute_reshard_plan(
         torch.cuda.synchronize()
         if service.requires_process_group_barrier:
             dist.barrier(group=group)
-        _refresh_module_caches(dst_module)
+        refresh_module_caches(dst_module)
         torch.cuda.synchronize()
         logger.info("Reshard complete")
         return
@@ -253,7 +263,7 @@ def get_sendable(param_name: str, param: torch.nn.Parameter) -> torch.Tensor:
             dst_param.quantize_(full_bf16)
     pending_quantized.clear()
 
-    _refresh_module_caches(dst_module)
+    refresh_module_caches(dst_module)
 
     # Ensure all writeback and cache-refresh copies are visible to subsequent CUDA
     # ops (e.g. CUDA graph warmup).  The synchronize() above fires *before* the

megatron/core/resharding/transforms.py

@@ -280,9 +280,7 @@ def finalize_recv(self, param_name, dst_slice, recv_buffers):
                         f"1D-scale param {param_name!r}: received {written} elements, "
                         f"expected {buf.data.numel()} (duplicate or missing slices?)"
                     )
-                mxfp8 = MXFP8Tensor.from_bf16(accum, backend=buf.backend)
-                buf.data.copy_(mxfp8.data)
-                buf.scale.view(torch.uint8).copy_(mxfp8.scale.view(torch.uint8))
+                buf.copy_(accum)
                 del self._pending_1d[buf_key]
             else:
                 self._pending_1d[buf_key][1] = written

megatron/core/transformer/moe/experts.py

@@ -1214,6 +1214,7 @@ def _build_concatenated_mxfp8_weights(self):
         for linear_name, buf_name in [('linear_fc1', '_fc1_weight'), ('linear_fc2', '_fc2_weight')]:
             linear = getattr(self, linear_name)
             q_list, s_list = [], []
+            source_dtype: torch.dtype | None = None
             for i in range(self.num_local_experts):
                 w = getattr(linear, f'weight{i}')
                 if isinstance(w, MXFP8Tensor):
@@ -1228,14 +1229,25 @@ def _build_concatenated_mxfp8_weights(self):
                 validate_mxfp8_tensor(
                     mxfp8, expected_backend=backend, tensor_name=f"{linear_name}.weight{i}"
                 )
+                if mxfp8.dtype is not None:
+                    source_dtype = source_dtype or mxfp8.dtype
+                    if mxfp8.dtype != source_dtype:
+                        raise RuntimeError(
+                            f"Conflicting source dtypes for {linear_name} expert weights: "
+                            f"{source_dtype} and {mxfp8.dtype}."
+                        )
                 q_list.append(mxfp8.data)
                 s_list.append(mxfp8.scale)
 
             stacked_data = torch.stack(q_list, dim=0).contiguous()
             stacked_scale = torch.stack(s_list, dim=0).contiguous()
 
             setattr(
-                self, buf_name, MXFP8Tensor(data=stacked_data, scale=stacked_scale, backend=backend)
+                self,
+                buf_name,
+                MXFP8Tensor(
+                    data=stacked_data, scale=stacked_scale, dtype=source_dtype, backend=backend
+                ),
             )
 
             # Redirect per-expert weight .data to views into the stacked buffer,

tests/unit_tests/inference/test_mxfp8_utils.py

@@ -11,6 +11,8 @@
 import pytest
 import torch
 
+from megatron.core.inference.quantization.mxfp8_tensor import HAVE_FLASHINFER, MXFP8Tensor
+
 pytestmark = [
     pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"),
     pytest.mark.internal,
@@ -269,6 +271,121 @@ def test_restore_uint8_scale_dtype_after_resharding(self):
         assert scale.dtype == torch.float8_e8m0fnu
         torch.testing.assert_close(scale.view(torch.uint8), scale_bytes)
 
+    @pytest.mark.parametrize("backend", ["triton", "flashinfer"])
+    def test_copy_preserves_storage_and_logical_metadata(self, backend):
+        if backend == "flashinfer" and not HAVE_FLASHINFER:
+            pytest.skip("FlashInfer not available")
+
+        source = torch.randn(16, 128, device="cuda", dtype=torch.bfloat16)
+        tensor = MXFP8Tensor.from_bf16(source, backend=backend)
+        data_ptr = tensor.data.data_ptr()
+        scale_ptr = tensor.scale.data_ptr()
+
+        assert tensor.shape == source.shape
+        assert tensor.dtype == source.dtype
+        assert tensor.device == source.device
+
+        updated = torch.randn_like(source)
+        assert tensor.copy_(updated) is tensor
+        expected = MXFP8Tensor.from_bf16(updated, backend=backend)
+
+        assert tensor.data.data_ptr() == data_ptr
+        assert tensor.scale.data_ptr() == scale_ptr
+        assert torch.equal(tensor.data, expected.data)
+        assert torch.equal(tensor.scale.view(torch.uint8), expected.scale.view(torch.uint8))
+
+    def test_copy_rejects_shape_mismatch(self):
+        tensor = MXFP8Tensor.from_bf16(
+            torch.randn(16, 128, device="cuda", dtype=torch.bfloat16), backend="triton"
+        )
+        with pytest.raises(ValueError, match="shape mismatch"):
+            tensor.copy_(torch.randn(32, 128, device="cuda", dtype=torch.bfloat16))
+
+    def test_copy_requires_backend(self):
+        source = torch.randn(16, 128, device="cuda", dtype=torch.bfloat16)
+        quantized = MXFP8Tensor.from_bf16(source, backend="triton")
+        backendless = MXFP8Tensor(data=quantized.data, scale=quantized.scale, dtype=source.dtype)
+        with pytest.raises(ValueError, match="without a quantization backend"):
+            backendless.copy_(source)
+
+    def test_copy_preserves_source_dtype_for_legacy_constructor(self):
+        initial = MXFP8Tensor.from_bf16(
+            torch.randn(16, 128, device="cuda", dtype=torch.bfloat16), backend="triton"
+        )
+        tensor = MXFP8Tensor(initial.data.clone(), initial.scale.clone(), "triton")
+        source = torch.randn(16, 128, device="cuda", dtype=torch.float16)
+        expected = MXFP8Tensor.from_bf16(source, backend="triton")
+
+        assert tensor.dtype is None
+        tensor.copy_(source)
+        assert tensor.dtype == source.dtype
+        assert torch.equal(tensor.data, expected.data)
+        assert torch.equal(tensor.scale, expected.scale)
+
+    def test_failed_legacy_copy_keeps_dtype_unknown(self):
+        initial = MXFP8Tensor.from_bf16(
+            torch.randn(16, 128, device="cuda", dtype=torch.bfloat16), backend="triton"
+        )
+        tensor = MXFP8Tensor(
+            initial.data.clone(), initial.scale.clone(), "unsupported"  # type: ignore[arg-type]
+        )
+
+        with pytest.raises(ValueError, match="Unknown MXFP8 quantization backend"):
+            tensor.copy_(torch.randn(16, 128, device="cuda", dtype=torch.float16))
+        assert tensor.dtype is None
+
+    def test_copy_does_not_downcast_update_to_logical_dtype(self):
+        source = torch.randn(16, 128, device="cuda", dtype=torch.float16)
+        tensor = MXFP8Tensor.from_bf16(source, backend="triton")
+        update = torch.linspace(
+            -100_000, 100_000, steps=16 * 128, device="cuda", dtype=torch.float32
+        ).reshape(16, 128)
+        expected = MXFP8Tensor.from_bf16(update, backend="triton")
+
+        assert tensor.dtype == torch.float16
+        tensor.copy_(update)
+        assert tensor.dtype == torch.float16
+        assert torch.equal(tensor.data, expected.data)
+        assert torch.equal(tensor.scale.view(torch.uint8), expected.scale.view(torch.uint8))
+
+    def test_copy_normalizes_flashinfer_fp32_input_to_bf16(self, monkeypatch):
+        tensor = MXFP8Tensor(
+            data=torch.empty((16, 128), device="cuda", dtype=torch.float8_e4m3fn),
+            scale=torch.empty(512, device="cuda", dtype=torch.uint8),
+            backend="flashinfer",
+        )
+        update = torch.randn(16, 128, device="cuda", dtype=torch.float32)
+        expected_data = torch.zeros_like(tensor.data)
+        expected_scale = torch.zeros_like(tensor.scale)
+        quantizer_input = {}
+
+        def fake_from_bf16(cls, value, group_size=32, backend="flashinfer"):
+            quantizer_input["dtype"] = value.dtype
+            return cls(data=expected_data, scale=expected_scale, backend=backend, dtype=value.dtype)
+
+        monkeypatch.setattr(MXFP8Tensor, "from_bf16", classmethod(fake_from_bf16))
+
+        tensor.copy_(update)
+
+        assert quantizer_input["dtype"] == torch.bfloat16
+        assert tensor.dtype == torch.bfloat16
+        assert torch.equal(tensor.data, expected_data)
+        assert torch.equal(tensor.scale, expected_scale)
+
+    def test_copy_rejects_stacked_storage(self):
+        tensor = MXFP8Tensor.from_bf16(
+            torch.randn(16, 128, device="cuda", dtype=torch.bfloat16), backend="triton"
+        )
+        stacked = MXFP8Tensor(
+            data=torch.stack([tensor.data, tensor.data]),
+            scale=torch.stack([tensor.scale, tensor.scale]),
+            dtype=tensor.dtype,
+            backend=tensor.backend,
+        )
+
+        with pytest.raises(ValueError, match="require 2D destination storage"):
+            stacked.copy_(torch.randn(2, 16, 128, device="cuda", dtype=torch.bfloat16))
+
     def test_reject_invalid_scale_dtype(self):
         from megatron.core.inference.quantization.mxfp8_tensor import ensure_mxfp8_scale_dtype
 
@@ -284,9 +401,9 @@ def test_validate_unswizzled_2d_scale_geometry(self):
         data = torch.empty((64, 128), dtype=torch.float8_e4m3fn, device="cuda")
         scale = torch.empty((64, 4), dtype=torch.uint8, device="cuda")
 
-        validate_mxfp8_tensor(MXFP8Tensor(data, scale, backend="flashinfer"))
+        validate_mxfp8_tensor(MXFP8Tensor(data, scale, dtype=torch.bfloat16, backend="flashinfer"))
 
-        invalid = MXFP8Tensor(data, scale[:63], backend="flashinfer")
+        invalid = MXFP8Tensor(data, scale[:63], dtype=torch.bfloat16, backend="flashinfer")
         with pytest.raises(ValueError, match="2D scale has shape"):
             validate_mxfp8_tensor(invalid)
 
@@ -300,7 +417,7 @@ def test_missing_backend_error_is_actionable(self):
         scale = torch.empty(512, dtype=torch.uint8, device="cuda")
 
         with pytest.raises(ValueError, match="backend= explicitly"):
-            validate_mxfp8_tensor(MXFP8Tensor(data, scale))
+            validate_mxfp8_tensor(MXFP8Tensor(data, scale, dtype=torch.bfloat16))
 
     @pytest.mark.parametrize("M,K", [(16, 128), (64, 256), (128, 2688)])
     def test_from_bf16_triton(self, M, K):
@@ -725,6 +842,7 @@ def stack_weights() -> MXFP8Tensor:
             return MXFP8Tensor(
                 data=torch.stack([weight.data for weight in per_expert]),
                 scale=torch.stack([weight.scale for weight in per_expert]),
+                dtype=torch.bfloat16,
                 backend="triton",
             )