triton.experimental.gluon.language.amd.cdna5.get_scaled_upcast_fp4_scale_layout

triton.experimental.gluon.language.amd.cdna5.get_scaled_upcast_fp4_scale_layout(src, scale_factor, elem_type, axis, _semantic=None)

Return the scale layout required by an FP4 scaled_upcast.

Returns the layout to use for the scale tensor required by a scaled_upcast based on the input tensor layout and the scale factor. Raises if there is no valid scale layout for the requested scale factor.

Parameters:

  • src – Packed FP4 source tensor (int8/uint8) whose layout is used to infer the scale layout.

  • scale_factor (int) – The scale factor, i.e. the number of elements along axis sharing a single scale. The scale tensor’s extent along axis is the unpacked extent of src divided by this.

  • elem_type – Output element type of the scaled_upcast (fp16 or bf16). This must match the elem_type passed to scaled_upcast.

  • axis (int) – Dimension along which packed FP4 values are unpacked and scales are shared.

Returns:

The distributed layout to use for the scale tensor.