Skip to content

Mistral batched decode: LLaMA decode path + parameterized RoPE theta (12x batched, full serving stack)#2

Open
mikepapadim wants to merge 1 commit into
feat/static-batched-decodefrom
feat/mistral-batched-decode
Open

Mistral batched decode: LLaMA decode path + parameterized RoPE theta (12x batched, full serving stack)#2
mikepapadim wants to merge 1 commit into
feat/static-batched-decodefrom
feat/mistral-batched-decode

Conversation

@mikepapadim

Copy link
Copy Markdown
Owner

Mistral batched decode (LLaMA decode path, no new kernels)

Extends the batched-decode engine (PR beehive-lab#129) to Mistral with a tiny generalization — Mistral
is LLaMA-arch (uses LlamaState + LlamaTornadoWeights, all accessors on the base
Configuration), so it runs the existing LLaMA decode path unchanged except:

  1. Parameterized RoPE theta in the decode + paged-decode rope kernels (config.ropeTheta();
    Mistral = 1e6, was hardcoded). Prefill ropes untouched; Llama-1B re-verified coherent.
  2. Generalized the decode layer graph (LlamaFP16LayersBatchDecodeMMA) to the base
    Configuration and dropped the LlamaConfiguration cast in the engine.

No new kernels. Any LLaMA-family model (RMSNorm + SwiGLU + GQA + RoPE, no QK-norm, no sliding
window) now drops in the same way; the full serving stack (continuous batching, PagedAttention,
prefix caching, on-device sampling) is model-agnostic and applies unchanged.

Verified — Mistral-7B-Instruct-v0.3 FP16, RTX 4090 (CUDA backend)

mode result
batched B=32 331 tok/s = 12× vs 27 tok/s single-stream; all 32 streams bit-exact, coherent
continuous + paged + prefix (128 requests) prefix-consistent, coherent; 85% slot utilization
paging 88 / 256 blocks used = ~11.6× less KV than full reservation
prefix caching 82.8% fewer prefill tokens (48-tok shared prompt)

Output sample (greedy): "The capital of France is Paris. It is one of the most famous cities in
the world, known for its art, fashion, cuisine, and historical landmarks such as the Eiffel Tower,
Louvre Museum, and Notre…"

Base branch: feat/static-batched-decode (PR beehive-lab#129). Tracked in roadmap issue beehive-lab#130.

…Configuration + parameterize RoPE theta; full serving stack works (12x batched, coherent, prefix-consistent)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant