Framework-specific API

• PyTorch
  • Linear
    • Linear.forward()
    • Linear.set_tensor_parallel_group()
  • GroupedLinear
    • GroupedLinear.forward()
    • GroupedLinear.set_tensor_parallel_group()
  • LayerNorm
  • RMSNorm
  • LayerNormLinear
    • LayerNormLinear.forward()
    • LayerNormLinear.set_tensor_parallel_group()
  • LayerNormMLP
    • LayerNormMLP.forward()
  • DotProductAttention
    • DotProductAttention.forward()
    • DotProductAttention.set_context_parallel_group()
  • MultiheadAttention
    • MultiheadAttention.forward()
    • MultiheadAttention.set_context_parallel_group()
    • MultiheadAttention.set_tensor_parallel_group()
  • TransformerLayer
    • TransformerLayer.forward()
    • TransformerLayer.set_context_parallel_group()
    • TransformerLayer.set_tensor_parallel_group()
  • CudaRNGStatesTracker
    • CudaRNGStatesTracker.add()
    • CudaRNGStatesTracker.fork()
    • CudaRNGStatesTracker.get_states()
    • CudaRNGStatesTracker.reset()
    • CudaRNGStatesTracker.set_states()
  • autocast()
  • quantized_model_init()
  • checkpoint()
  • make_graphed_callables()
  • get_cpu_offload_context()
  • mark_not_offload()
  • ManualOffloadSynchronizer
  • parallel_cross_entropy()
  • Recipe availability
    • is_fp8_available()
    • is_mxfp8_available()
    • is_fp8_block_scaling_available()
    • is_nvfp4_available()
    • is_bf16_available()
    • get_cudnn_version()
    • get_device_compute_capability()
    • get_default_recipe()
  • Mixture of Experts (MoE) functions
    • moe_permute()
    • moe_permute_with_probs()
    • moe_unpermute()
    • moe_sort_chunks_by_index()
    • moe_sort_chunks_by_index_with_probs()
  • Communication-computation overlap
    • initialize_ub()
    • destroy_ub()
    • UserBufferQuantizationMode
      • UserBufferQuantizationMode.FP8
      • UserBufferQuantizationMode.NONE
  • Quantized tensors
    • QuantizedTensorStorage
      • QuantizedTensorStorage.prepare_for_saving()
      • QuantizedTensorStorage.restore_from_saved()
      • QuantizedTensorStorage.update_usage()
    • QuantizedTensor
      • QuantizedTensor.dequantize()
      • QuantizedTensor.quantize_()
    • Float8TensorStorage
    • MXFP8TensorStorage
    • Float8BlockwiseQTensorStorage
    • NVFP4TensorStorage
    • Float8Tensor
    • MXFP8Tensor
    • Float8BlockwiseQTensor
    • NVFP4Tensor
  • Quantizers
    • Quantizer
      • Quantizer.quantize()
      • Quantizer.update_quantized()
    • Float8Quantizer
    • Float8CurrentScalingQuantizer
    • MXFP8Quantizer
    • Float8BlockQuantizer
    • NVFP4Quantizer
  • Tensor saving and restoring functions
    • prepare_for_saving()
    • restore_from_saved()
  • Operation fuser
    • Sequential
      • Sequential.forward()
    • FusibleOperation
      • FusibleOperation.fuser_backward()
      • FusibleOperation.fuser_forward()
    • BasicOperation
      • BasicOperation.op_backward()
      • BasicOperation.op_forward()
    • FusedOperation
      • FusedOperation.fuser_backward()
      • FusedOperation.fuser_forward()
    • register_forward_fusion()
    • register_backward_fusion()
    • Linear
    • AddExtraInput
    • AllGather
    • AllReduce
    • BasicLinear
      • BasicLinear._functional_backward()
      • BasicLinear._functional_forward()
    • Bias
    • ClampedSwiGLU
    • ConstantScale
    • Dropout
    • GEGLU
    • GELU
    • GLU
    • GroupedLinear
    • Identity
    • L2Normalization
    • LayerNorm
    • MakeExtraOutput
    • QGELU
    • QGEGLU
    • Quantize
    • ReGLU
    • ReLU
    • ReduceScatter
    • Reshape
    • RMSNorm
    • SReGLU
    • SReLU
    • ScaledSwiGLU
    • SiLU
    • SwiGLU
  • Deprecated functions
    • fp8_autocast()
    • fp8_model_init()
• Jax
  • Pre-defined Variable of Logical Axes
  • Checkpointing
  • Modules
    • TransformerLayerType
    • MeshResource
      • MeshResource.dp_resource
      • MeshResource.tp_resource
      • MeshResource.tpsp_resource
      • MeshResource.fsdp_resource
      • MeshResource.pp_resource
      • MeshResource.cp_resource
    • fp8_autocast()
    • autocast()
    • update_collections()
    • LayerNorm
      • LayerNorm.__call__()
    • DenseGeneral
      • DenseGeneral.__call__()
    • LayerNormDenseGeneral
      • LayerNormDenseGeneral.__call__()
    • LayerNormMLP
      • LayerNormMLP.__call__()
    • RelativePositionBiases
      • RelativePositionBiases.__call__()
    • DotProductAttention
      • DotProductAttention.__call__()
    • MultiHeadAttention
      • MultiHeadAttention.__call__()
    • TransformerLayer
      • TransformerLayer.__call__()
    • extend_logical_axis_rules()