Framework-specific API

• PyTorch
  • Linear
    • forward()
  • GroupedLinear
    • forward()
  • LayerNorm
  • RMSNorm
  • LayerNormLinear
    • forward()
  • LayerNormMLP
    • forward()
  • DotProductAttention
    • forward()
    • set_context_parallel_group()
  • MultiheadAttention
    • forward()
    • set_context_parallel_group()
    • set_tensor_parallel_group()
  • TransformerLayer
    • forward()
    • set_context_parallel_group()
    • set_tensor_parallel_group()
  • CudaRNGStatesTracker
    • add()
    • fork()
    • get_states()
    • reset()
    • set_states()
  • autocast
  • quantized_model_init()
  • checkpoint()
  • make_graphed_callables()
  • get_cpu_offload_context()
  • mark_not_offload()
  • ManualOffloadSynchronizer
  • parallel_cross_entropy()
  • interleave_glu_tensor()
  • deinterleave_glu_tensor()
  • Data types
    • DType
      • kBFloat16
      • kByte
      • kFloat16
      • kFloat32
      • kFloat4E2M1
      • kFloat8E4M3
      • kFloat8E5M2
      • kInt32
  • 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
      • FP8
      • NONE
  • Quantized tensors
    • QuantizedTensorStorage
      • prepare_for_saving()
      • restore_from_saved()
      • update_usage()
    • QuantizedTensor
      • dequantize()
      • quantize_()
    • Float8TensorStorage
    • MXFP8TensorStorage
    • Float8BlockwiseQTensorStorage
    • NVFP4TensorStorage
    • Float8Tensor
    • MXFP8Tensor
    • Float8BlockwiseQTensor
    • NVFP4Tensor
  • Quantizers
    • Quantizer
      • quantize()
      • update_quantized()
    • Float8Quantizer
    • Float8CurrentScalingQuantizer
    • MXFP8Quantizer
    • Float8BlockQuantizer
    • NVFP4Quantizer
  • Tensor saving and restoring functions
    • prepare_for_saving()
    • restore_from_saved()
  • Operation fuser
    • Sequential
      • forward()
    • FusibleOperation
      • fuser_backward()
      • fuser_forward()
    • BasicOperation
      • op_backward()
      • op_forward()
    • FusedOperation
      • fuser_backward()
      • fuser_forward()
    • register_forward_fusion()
    • register_backward_fusion()
    • register_forward_backward_fusion()
    • Linear
    • AddExtraInput
    • AllGather
    • AllReduce
    • BasicLinear
      • _functional_backward()
      • _functional_forward()
    • Bias
    • ClampedSwiGLU
    • ConstantScale
    • Dropout
    • GEGLU
    • GELU
    • GLU
    • GroupedLinear
    • Identity
    • L2Normalization
    • LayerNorm
    • MakeExtraOutput
    • QGELU
    • QGEGLU
    • Quantize
    • ReGLU
    • ReLU
    • ReduceScatter
    • Reshape
    • RMSNorm
    • SReGLU
    • SReLU
    • ScaledClampedQGeGLU
    • ScaledSwiGLU
    • SiLU
    • SwiGLU
    • mhc_fused_sinkhorn()
    • mhc_fused_scale()
    • mhc_fused_aggregate()
    • mhc_fused_expand_combine()
    • mhc_fused_projection()
  • Deprecated functions
    • fp8_autocast()
    • fp8_model_init()
• Jax
  • Pre-defined Variable of Logical Axes
  • Checkpointing
  • Modules
    • TransformerLayerType
    • MeshResource
      • dp_resource
      • tp_resource
      • tpsp_resource
      • fsdp_resource
      • pp_resource
      • cp_resource
      • ep_resource
    • fp8_autocast()
    • autocast()
    • update_collections()
    • LayerNorm
      • __call__()
    • DenseGeneral
      • __call__()
    • LayerNormDenseGeneral
      • __call__()
    • LayerNormMLP
      • __call__()
    • RelativePositionBiases
      • __call__()
    • DotProductAttention
      • __call__()
    • MultiHeadAttention
      • __call__()
    • TransformerLayer
      • __call__()
    • extend_logical_axis_rules()