huggingface/candle · error

softmax-last-dim is not implemented for {dtype:?}

Error message

softmax-last-dim is not implemented for {dtype:?}

What it means

The Metal softmax kernel dispatches to named metal shaders softmax_f32/softmax_f16/softmax_bf16 only. Any other dtype (F64, integers, etc.) has no GPU kernel registered, so metal_fwd bails with this message naming the dtype.

Source

Thrown at candle-nn/src/ops.rs:404

        Ok((dst, layout.shape().clone()))
    }

    #[cfg(feature = "metal")]
    fn metal_fwd(
        &self,
        storage: &candle::MetalStorage,
        layout: &Layout,
    ) -> Result<(candle::MetalStorage, Shape)> {
        use candle::backend::BackendStorage;
        let device = storage.device();
        let encoder = device.command_encoder()?;
        encoder.set_label("softmax");
        let kernels = device.kernels();
        let name = match storage.dtype() {
            DType::F32 => "softmax_f32",
            DType::F16 => "softmax_f16",
            DType::BF16 => "softmax_bf16",
            dtype => candle::bail!("softmax-last-dim is not implemented for {dtype:?}"),
        };

        let n = layout.stride().len();
        if !(layout.is_contiguous() && layout.stride()[n - 1] == 1) {
            candle::bail!("Non contiguous softmax-last-dim is not implemented");
        }

        let last_dim = layout.dims()[layout.shape().rank() - 1];
        let elem_count = layout.shape().elem_count();
        let output = device
            .new_buffer_builder()
            .with_size_for(elem_count, storage.dtype())
            .with_label("softmax")
            .build()?;
        candle_metal_kernels::call_last_softmax(
            device.metal_device(),
            &encoder,
            kernels,

View on GitHub (pinned to d5fee525bf)

Solutions

  1. Cast to F32 (or F16/BF16) with .to_dtype() before the softmax
  2. Keep model tensors in F32/F16/BF16 on Metal; avoid F64 on GPU backends
  3. Verify the op's device matches the intended backend; CPU supports F64 softmax but Metal does not

Example fix

// before
let probs = softmax_last_dim(&logits.to_device(&metal_device)?)?; // F64
// after
let probs = softmax_last_dim(&logits.to_dtype(DType::F32)?.to_device(&metal_device)?)?;
Defensive patterns

Strategy: validation

Validate before calling

fn ensure_metal_softmax(t: &Tensor) -> Result<()> {
    if !matches!(t.dtype(), DType::F32 | DType::F16 | DType::BF16) {
        bail!("Metal softmax supports F32/F16/BF16 only, got {:?}", t.dtype());
    }
    Ok(())
}

Type guard

fn metal_softmax_ok(d: DType) -> bool { matches!(d, DType::F32 | DType::F16 | DType::BF16) }

Try / catch

let t = if metal_softmax_ok(t.dtype()) { t } else { t.to_dtype(DType::F32)? };

Prevention

When it happens

Trigger: Calling ops::softmax_last_dim on a Metal-device tensor with a dtype other than F32/F16/BF16, such as F64 or I64.

Common situations: Default F64 tensors created from Rust f64 arrays and sent to a Metal device; index/integer tensors accidentally passed to softmax; dtype changed upstream by a model config.

Related errors

AI-assisted analysis of huggingface/candle@d5fee525bf (2026-09-02). Data as JSON: /api/errors/8514524bf257c7d3. Report an issue: GitHub.