Greptile Summary

This PR restores the get_skip_fp8_weight_update_tensor and set_skip_fp8_weight_update_tensor classmethods on FP8GlobalStateManager that were removed in PR #2759, which broke external consumers such as Megatron-LM's CUDA graph implementation.

• quantization.py: Adds the two classmethods; the setter lazily initializes the backing CUDA float32 tensor on first call, making it also safe against the previously possible None.fill_() crash path.
• graph.py: Replaces two sites of direct quantization_state.skip_fp8_weight_update_tensor attribute access with the new setter, improving encapsulation and gaining the null-safety benefit for free.

Confidence Score: 5/5

Safe to merge — the change is a narrow, focused restoration of a deleted public interface with no behavioral regressions.

The two new classmethods are simple wrappers around already-tested state; the setter's lazy initialization is an improvement over the previous bare .fill_() call that could have panicked if the tensor was None. Internal module files continue to access the attribute directly, which is consistent with how they worked before and is not broken by this change.

No files require special attention.

Important Files Changed

Sequence Diagram

sequenceDiagram
    participant MCore as MCore (cuda_graphs.py)
    participant FGSM as FP8GlobalStateManager
    participant State as FP8GlobalState

    Note over MCore,State: Setter — called during CUDA graph capture setup
    MCore->>FGSM: set_skip_fp8_weight_update_tensor(True/False)
    FGSM->>State: "create tensor if None (device=cuda, dtype=float32)"
    FGSM->>State: fill_(skip)

    Note over MCore,State: Getter — called to read current flag
    MCore->>FGSM: get_skip_fp8_weight_update_tensor()
    FGSM->>State: skip_fp8_weight_update_tensor
    State-->>MCore: Optional[torch.Tensor]

_{Reviews (4): Last reviewed commit: "Merge branch 'main' into xren/fix_skip_f..." | Re-trigger Greptile}

/te-ci pytorch