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()
    • LayerNormMLP.set_tensor_parallel_group()
  • 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()
  • 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()
  • 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()