Purpose

• FP4
  • Fix bug discovered here where dynamic="local" nvfp4 calculations would increment the observer twice as fast as normal
  • Enable MSE observer to be used with FP4
• Simplification
  • Make supporting attention calibration easier by separating out weight/activation/attention reshaping
  • Improve readability of observer codes by removing many levels of function indirection
  • Drop support for calibration with non-divisible group sizes. This is not really a loss, since forward passes also make this assumption

Before	After

Changes

• Standardize reshaping using flatten_for_calibration
  • This function reshapes all observed values to (num_observations, *qparams_shape, group_size)
  • This function the complexity associated with passing "reduce dims" and trying to handle weights, activations, and attention states all in the same function
  • In the future, this function could be applied to the quantization forward pass, although there's probably no need to outside of standardization
• Implement get_global_scale on Observer base
  • This function decouples minmax calculations from regular qparam calculations (avoiding the double increment bug)
  • This function enables the MSE observer to be used with FP4 global scales

Testing

• Added additional minmax tests which check exact values of scales. This test passes both on main and this branch, demonstrating that minmax observer behavior remains unchanged
• Added additional MSE tests which check exact values of mse losses. This test passes both on main and this branch, demonstrating that MSE observer behavior remains unchanged