fix: auto-adjust n_ubatch when n_batch > 512 for embedding

[lxcxjxhx]

Motivation

Fixes #1762

When using Llama.embed() with n_batch > 512, the kernel crashes because n_ubatch defaults to 512. Users had to manually set both n_batch and n_ubatch to the same value as a workaround:

embedder = Llama.from_pretrained(
    repo_id="lm-kit/bge-m3-gguf",
    filename="*F16.gguf",
    n_ctx=0,
    n_gpu_layers=-1,
    n_batch=4096,
    n_ubatch=4096,  # Manual workaround required
    embedding=True,
    pooling_type=LLAMA_POOLING_TYPE_NONE,
    verbose=False
)

This is unintuitive and causes crashes when users only set n_batch without realizing they also need to set n_ubatch.

Changes

Modified llama_cpp/llama.py in the embed() method:

Before:

def embed(self, input, normalize=False, truncate=True, return_count=False):
    n_embd = self.n_embd()
    n_batch = self.n_batch
    n_seq_max = self.context_params.n_seq_max
    # ... rest of method

After:

def embed(self, input, normalize=False, truncate=True, return_count=False):
    n_embd = self.n_embd()
    n_batch = self.n_batch
    n_seq_max = self.context_params.n_seq_max
    
    # Fix for issue #1762: ensure n_ubatch is at least n_batch for embedding
    # When n_batch > 512, n_ubatch must also be increased to avoid crashes
    if self.context_params.n_ubatch < n_batch:
        self.set_n_ubatch(n_batch)
    
    # ... rest of method

This change:

1. Automatically adjusts n_ubatch to match n_batch when needed
2. Prevents crashes when using n_batch > 512
3. Maintains backward compatibility - existing code with default n_batch=512 is unaffected
4. Provides a more intuitive API for long-context embedding models like BGE-M3

Testing

Added test file tests/test_embed_n_ubatch_fix.py with:

• Test verifying n_ubatch is auto-adjusted when n_batch > n_ubatch
• Test verifying no adjustment when n_ubatch is already sufficient
• Verified embedding works with n_batch=4096 without manual n_ubatch
• Confirmed backward compatibility with default n_batch=512

Impact

• Enables embedding with n_batch > 512 without manual configuration
• Fixes crashes for long-context embedding models (BGE-M3, etc.)
• No breaking changes to existing API
• Improves user experience for late interaction retrieval use cases

Checklist

• Fixes Llama.embed crashes when n_batch > 512 #1762
• Code changes are minimal and focused on the issue
• Backward compatibility maintained
• Test cases added
• No breaking changes to existing API
• Follows existing code style and patterns