For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).
Mojo struct
BundledAllReduceAddRMSNormQuantFP8
struct BundledAllReduceAddRMSNormQuantFP8
Registers the mo.composite.bundled.allreduce_add_rms_norm_quant_fp8 graph op with the graph compiler.
Implemented traits
Methods
execute
static def execute[dtype: DType, output_type: DType, scales_type: DType, rank: Int, target: StringSlice[ImmStaticOrigin], _trace_name: StringSlice[ImmStaticOrigin]](output: ManagedTensorSlice[IOSpec[_, _].Output, static_spec=output.static_spec], out_scale: ManagedTensorSlice[IOSpec[_, _].Output, static_spec=out_scale.static_spec], out_residual: ManagedTensorSlice[IOSpec[_, _].Output, static_spec=out_residual.static_spec], inputs: VariadicTensors[IOSpec[_, _].Input, static_specs=inputs.static_specs], signal_buffers: VariadicTensors[IOSpec[_, _].MutableInput, static_specs=signal_buffers.static_specs], residual: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=residual.static_spec], gamma: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=gamma.static_spec], epsilon: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=epsilon.static_spec], weight_offset: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=weight_offset.static_spec], scale_ub: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=scale_ub.static_spec], ctx: DeviceContext)
Per-device fused allreduce.sum + add + rms_norm + fp8 quantize.
Single-device analog of DistributedAllReduceAddRMSNormQuantFP8, for
use inside mo.parallel. The parallel framework launches one
instance per GPU; this kernel invokes the same underlying primitive
(allreduce_residual_rmsnorm) that the distributed variant calls
from within _launch_device_collective, but for a single device.
Parameters:
- dtype (
DType): Element type of the input and residual tensors. - output_type (
DType): Element type of the quantized output tensor. - scales_type (
DType): Element type of the per-token scale tensor. - rank (
Int): Number of dimensions in the input and output tensors. - target (
StringSlice[ImmStaticOrigin]): Target device identifier for code generation. - _trace_name (
StringSlice[ImmStaticOrigin]): Trace name used for profiling and debugging.
Args:
- output (
ManagedTensorSlice[IOSpec[_, _].Output, static_spec=output.static_spec]): FP8 quantized output tensor for THIS GPU. - out_scale (
ManagedTensorSlice[IOSpec[_, _].Output, static_spec=out_scale.static_spec]): Per-token scale tensor for THIS GPU. - out_residual (
ManagedTensorSlice[IOSpec[_, _].Output, static_spec=out_residual.static_spec]): Post-add residual tensor for THIS GPU. - inputs (
VariadicTensors[IOSpec[_, _].Input, static_specs=inputs.static_specs]): Input tensors from ALL participating GPUs. - signal_buffers (
VariadicTensors[IOSpec[_, _].MutableInput, static_specs=signal_buffers.static_specs]): Signal buffers for ALL participating GPUs. - residual (
ManagedTensorSlice[IOSpec[_, _].Input, static_spec=residual.static_spec]): Residual tensor for THIS GPU. - gamma (
ManagedTensorSlice[IOSpec[_, _].Input, static_spec=gamma.static_spec]): RMSNorm weight for THIS GPU. - epsilon (
ManagedTensorSlice[IOSpec[_, _].Input, static_spec=epsilon.static_spec]): RMSNorm epsilon scalar (host). - weight_offset (
ManagedTensorSlice[IOSpec[_, _].Input, static_spec=weight_offset.static_spec]): RMSNorm weight offset scalar (host). - scale_ub (
ManagedTensorSlice[IOSpec[_, _].Input, static_spec=scale_ub.static_spec]): Quantization scale upper bound scalar (host). - ctx (
DeviceContext): Device context for THIS GPU.