Skip to content

Optimize implementation, add multi stream support - #6

Merged
timkhronos merged 5 commits into
MSAfrom
10%Decode-speed,-10%-PP-speed,-20%-smalle-compute-buffer,-F16-indexer-cache-and-Multi-stream-support
Jul 11, 2026
Merged

Optimize implementation, add multi stream support#6
timkhronos merged 5 commits into
MSAfrom
10%Decode-speed,-10%-PP-speed,-20%-smalle-compute-buffer,-F16-indexer-cache-and-Multi-stream-support

Conversation

@timkhronos

Copy link
Copy Markdown
Owner

Overview

Fully rewrote minimax-m3.cpp for speed and buffer size gains:

Unified the 4-way + decode, 1 FA call per layer instead of 4, with the groups mapped onto ne[3]

Custom CPU op now emits block-level mask, expanded on GPU, which causes CPU to GPU transfer to shrinks at prefill

Decode: ~25 nodes/layer vs ~50, no per-group concats/conts

Unified selection semantics, so both regimes rank bs + local bias (position-anchored local force), which means prefill/decode can no longer disagree on selection

can_reuse on the MSA bias input. Graph reuse at decode restored (was rebuilding the full graph every token)

In-place mask adds, shrinking compute buffer ~6.8 to ~4.2 GiB at ub2048/62k

Multi-stream: MSA now runs with -np N when kv_unified=false. Decode stays batched across streams (still 1 FA call), prefill loops per stream. dense fallback only for --kv-unified + multi-seq

Additional information

src/llama-kv-cache.cpp: k_idx cache f32 -> f16

src/models/models.h: two MSA method decls, instead one build_attn_msa_fa(…, Gp, …)

Measured effect: decode 6.2(4WAY)–7.15(MSA_decode) -> 7.7~7.8 t/s, flat from 5k to 60k+. prefill around 10% faster. buffer about 20% smaller, multi-user support.

Requirements

@github-actions github-actions Bot added the model label Jul 11, 2026
@timkhronos
timkhronos merged commit cacc42f into MSA Jul 11, 2026
4 of 23 checks passed
@timkhronos
timkhronos deleted the 10%Decode-speed,-10%-PP-speed,-20%-smalle-compute-buffer,-F16-indexer-cache-and-Multi-stream-support branch July 29, 2026 19:09
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant