Add more misc. changes from candle fork

[EricLBuehler]

• indexed_moe_forward (fast path for ggml quants)
• Improved usability of Context
• Add full attn support for Metal SDPA
• Fix bug w/ FlashAttn f16
• Add necessary metal Device apis

[ivarflakstad]

Excellent stuff

[ivarflakstad]

Just thinking out loud here. It would be nice to have automatic fallback to an approach that isn't as optimized, but still valid. Perhaps returning Result<Option<(CudaStorage, crate::Shape)>> is a decent starting point?
If None then fallback?

Not thinking we add this in this PR ofc.

[EricLBuehler]

This might work, the issue is that effectively indexed_moe_forward is a grouped gemm so we'd need existing infrastructure to run a grouped gemm.

Regardless, providing a grouped gemm functionality will be very useful!

[EricLBuehler]

Addressed the review comments, the new_private_buffer method is now implemented correctly.

[haricot]

for ci ubuntu, the linker seems to have crashed due to lack of memory.
maybe possible resolution:

1. workflow .github/workflows/rust-ci.yml:

      # Add lld install, optional if lld is already present on runners
      - name: Install lld (Linux only)
        if: runner.os == 'Linux'
        run: sudo apt-get update && sudo apt-get install -y lld

      # The change: Add RUSTFLAGS for Linux to use linker-features
      - name: Run tests (with lld on Linux)
        if: runner.os == 'Linux'
        env:
          RUSTFLAGS: "-Clinker-features=-lld"
        run: cargo test --workspace

      # Existing Windows and Mac steps (unchanged)
      - name: Run tests (Windows & macOS)
        if: runner.os != 'Linux'
        run: cargo test --workspace

2. add : .cargo/config.toml
   [target.x86_64-unknown-linux-gnu] rustflags = ["-Clinker-features=-lld"]

[ivarflakstad]

Lgtm! 🔥
Same wrt CI here as well

[ivarflakstad]

@haricot yeah we can try setting the flag in CI.
I'd prefer to not have candle be opinionated wrt the linker, so I'd rather not add it to the config if we can avoid it.
I assume there are some (possibly obscure) reasons why it is not the default.