From 280aabd93b26e1707e54e8d0cca139236b92c5ac Mon Sep 17 00:00:00 2001 From: octo-patch Date: Tue, 2 Jun 2026 05:08:09 +0800 Subject: [PATCH] feat: upgrade MiniMax default validation target to M3 - Add MiniMax-M3 as the new default target for validate_minimax_m2.py - Keep MiniMax-M2.7 available via --model MiniMax-M2.7 (legacy fallback) - Refactor MODEL_NAME constant to SUPPORTED_MODELS / DEFAULT_MODEL pair - Update README compatibility section to highlight M3 as new flagship while documenting M2.7 as fallback - Update test docstrings to reflect M3 / M2.7 dual targeting The architecture constants (head_dim=128, num_kv_heads=8, num_layers=62) remain unchanged because the MiniMax-M3 MoE attention layout matches MiniMax-M2.7, so existing IsoQuant / PlanarQuant block-rotation tests continue to validate both checkpoints with no code changes. --- README.md | 18 ++++++---- tests/test_minimax_m2.py | 29 ++++++++-------- turboquant/validate_minimax_m2.py | 55 ++++++++++++++++++------------- 3 files changed, 58 insertions(+), 44 deletions(-) diff --git a/README.md b/README.md index 6797f12..3b0ce0c 100644 --- a/README.md +++ b/README.md @@ -182,12 +182,15 @@ python -m turboquant.benchmark_triton # Triton kernel speed python -m turboquant.poc_high_context --backend planar # High-context generation ``` -## MiniMax-M2.7 Compatibility +## MiniMax-M3 / M2.7 Compatibility -IsoQuant and PlanarQuant work out-of-the-box with [MiniMax-M2.7](https://huggingface.co/MiniMaxAI/MiniMax-M2.7), -a 230B MoE model with 204K context and Grouped Query Attention (48 query heads / 8 KV heads). +IsoQuant and PlanarQuant work out-of-the-box with the [MiniMax](https://huggingface.co/MiniMaxAI) +family. The new flagship [MiniMax-M3](https://huggingface.co/MiniMaxAI/MiniMax-M3) is the +default validation target; [MiniMax-M2.7](https://huggingface.co/MiniMaxAI/MiniMax-M2.7) +remains supported as a fallback. Both are 230B MoE models with 204K context and Grouped +Query Attention (48 query heads / 8 KV heads). -**Architecture fit:** +**Architecture fit (M3 and M2.7 share the same MoE attention layout):** - `head_dim = 128` — perfect alignment for IsoQuant 4D blocks (128 / 4 = 32 groups) and PlanarQuant 2D pairs (128 / 2 = 64 groups) - `num_kv_heads = 8` — each head compressed independently, GQA expansion handled by the model @@ -202,9 +205,10 @@ a 230B MoE model with 204K context and Grouped Query Attention (48 query heads / **Validate (requires GPU + `pip install -e ".[validate]"`):** ```bash -python -m turboquant.validate_minimax_m2 # synthetic + real model -python -m turboquant.validate_minimax_m2 --dry-run # synthetic only (no model download) -python -m pytest tests/test_minimax_m2.py -v # unit tests (no model required) +python -m turboquant.validate_minimax_m2 # M3 (default) — synthetic + real model +python -m turboquant.validate_minimax_m2 --model MiniMax-M2.7 # legacy M2.7 target +python -m turboquant.validate_minimax_m2 --dry-run # synthetic only (no model download) +python -m pytest tests/test_minimax_m2.py -v # unit tests (no model required) ``` **[ParaMind2025](https://github.com/ParaMind2025)** — PlanarQuant (2D Givens rotation) and IsoQuant (4D quaternion rotation) were designed by ParaMind2025. Their insight that simple block-diagonal rotations could match full-rank transforms for KV cache decorrelation made the llama.cpp integration practical. diff --git a/tests/test_minimax_m2.py b/tests/test_minimax_m2.py index c86896e..68816f8 100644 --- a/tests/test_minimax_m2.py +++ b/tests/test_minimax_m2.py @@ -1,8 +1,9 @@ """ -Tests for MiniMax-M2.7 KV cache compression compatibility. +Tests for MiniMax-M3 / M2.7 KV cache compression compatibility. Validates that IsoQuant / PlanarQuant / LiteratiQuant work correctly -with MiniMax-M2.7's architecture dimensions: +with the MiniMax family architecture dimensions (M3 is the new default, +M2.7 is the previous flagship and shares the same MoE attention layout): - head_dim = 128 - num_kv_heads = 8 (GQA: 48 query heads, 8 KV heads) - num_layers = 62 @@ -24,7 +25,7 @@ from turboquant.literatiquant import LiteratiQuantMSE, LiteratiQuantKVCache from turboquant.compressors import TurboQuantCompressorV2, TurboQuantCompressorMSE -# ── MiniMax-M2.7 architecture constants ────────────────────────────── +# ── MiniMax-M3 / M2.7 architecture constants ───────────────────────── HEAD_DIM = 128 NUM_KV_HEADS = 8 NUM_ATTN_HEADS = 48 @@ -36,7 +37,7 @@ @pytest.fixture(scope="module") def minimax_kv_cache(): - """Synthetic KV cache with MiniMax-M2.7 dimensions (single-batch).""" + """Synthetic KV cache with MiniMax-M3 / M2.7 dimensions (single-batch).""" torch.manual_seed(42) seq_len = 512 batch = 1 @@ -49,7 +50,7 @@ def minimax_kv_cache(): @pytest.fixture(scope="module") def minimax_query(): - """Synthetic query tensor matching MiniMax-M2.7 attention head dim.""" + """Synthetic query tensor matching MiniMax attention head dim.""" torch.manual_seed(99) # Queries use full num_attn_heads; KV use num_kv_heads (GQA) q = torch.randn(1, NUM_ATTN_HEADS, 1, HEAD_DIM) @@ -59,7 +60,7 @@ def minimax_query(): # ── IsoQuantMSE tests ───────────────────────────────────────────────── class TestIsoQuantMiniMax: - """IsoQuant on MiniMax-M2.7 head_dim=128 vectors.""" + """IsoQuant on MiniMax-M3 / M2.7 head_dim=128 vectors.""" @pytest.mark.parametrize("bits", [2, 3, 4]) def test_output_shape(self, minimax_kv_cache, bits): @@ -83,7 +84,7 @@ def test_mse_within_bounds(self, minimax_kv_cache, bits): assert mse < 1.5, f"MSE {mse:.4f} unexpectedly high for {bits}-bit IsoQuant" def test_mse_decreases_with_bits(self, minimax_kv_cache): - """Higher bits → lower MSE on MiniMax-M2.7 dimensions.""" + """Higher bits → lower MSE on MiniMax head dimensions.""" keys, _ = minimax_kv_cache flat_keys = keys.view(-1, HEAD_DIM) mses = [] @@ -106,7 +107,7 @@ def test_head_dim_128_alignment(self): # ── PlanarQuantMSE tests ────────────────────────────────────────────── class TestPlanarQuantMiniMax: - """PlanarQuant on MiniMax-M2.7 dimensions.""" + """PlanarQuant on MiniMax-M3 / M2.7 dimensions.""" @pytest.mark.parametrize("bits", [2, 3, 4]) def test_output_shape(self, minimax_kv_cache, bits): @@ -140,7 +141,7 @@ class TestGQACompatibility: """Tests that the quantizers handle GQA (num_kv_heads != num_query_heads).""" def test_kv_head_count(self): - """MiniMax-M2.7 GQA: 8 KV heads serving 48 query heads.""" + """MiniMax M3 / M2.7 GQA: 8 KV heads serving 48 query heads.""" assert NUM_KV_HEADS == 8 assert NUM_ATTN_HEADS == 48 assert NUM_ATTN_HEADS % NUM_KV_HEADS == 0 # clean group ratio @@ -178,7 +179,7 @@ def test_batch_compress_all_kv_heads(self, minimax_kv_cache): class TestAttentionScoreFidelity: """ Verify that compressed KV cache preserves top-k attention token ranking, - which is critical for MiniMax-M2.7's long-context capabilities (204K). + which is critical for MiniMax's long-context capabilities (204K). """ @pytest.mark.parametrize("bits", [3, 4]) @@ -236,7 +237,7 @@ def test_score_cosine_similarity(self, minimax_kv_cache, bits): avg_cos = sum(cos_sims) / len(cos_sims) assert avg_cos > 0.8, ( f"Average cosine sim {avg_cos:.4f} too low for {bits}-bit " - f"IsoQuant on MiniMax-M2.7 dimensions" + f"IsoQuant on MiniMax head dimensions" ) @@ -246,7 +247,7 @@ class TestLiteratiQuantMiniMax: """1-bit KV cache compression for ultra-high compression ratios.""" def test_kv_cache_compress_decompress(self, minimax_kv_cache): - """LiteratiQuantKVCache round-trip on MiniMax-M2.7 KV tensors.""" + """LiteratiQuantKVCache round-trip on MiniMax KV tensors.""" keys, _ = minimax_kv_cache B, H, S, D = keys.shape @@ -274,7 +275,7 @@ def test_group_size_alignment(self): # ── TurboQuantCompressor tests ──────────────────────────────────────── class TestTurboQuantMiniMax: - """TurboQuant asymmetric estimator with MiniMax-M2.7 dimensions.""" + """TurboQuant asymmetric estimator with MiniMax-M3 / M2.7 dimensions.""" @pytest.mark.parametrize("bits", [2, 3, 4]) def test_compress_and_reconstruct_shape(self, minimax_kv_cache, bits): @@ -312,7 +313,7 @@ def test_inner_product_unbiased(self, minimax_kv_cache, bits): # ── Memory estimate tests ───────────────────────────────────────────── class TestMemoryEstimate: - """Verify compression ratio math for MiniMax-M2.7 long-context scenarios.""" + """Verify compression ratio math for MiniMax long-context scenarios.""" @pytest.mark.parametrize("seq_len,bits", [ (32_768, 3), diff --git a/turboquant/validate_minimax_m2.py b/turboquant/validate_minimax_m2.py index fb7c703..89f0d23 100644 --- a/turboquant/validate_minimax_m2.py +++ b/turboquant/validate_minimax_m2.py @@ -1,8 +1,9 @@ """ -MiniMax-M2.7 KV Cache Compression Validation. +MiniMax-M3 / M2.7 KV Cache Compression Validation. Demonstrates IsoQuant / PlanarQuant KV cache compression on -MiniMax-M2.7 (MiniMaxAI/MiniMax-M2.7), a 230B MoE model with: +the MiniMax family (MiniMaxAI/MiniMax-M3 by default, MiniMax-M2.7 still +supported as a fallback). Both share the same MoE attention layout: - head_dim = 128 (4D quaternion blocks align perfectly) - num_kv_heads = 8 (GQA: 48 query heads, 8 KV heads) - num_layers = 62 @@ -10,12 +11,13 @@ Requirements: pip install -e ".[validate]" - # MiniMax-M2.7 requires trust_remote_code=True and ~48 GB GPU VRAM + # The MiniMax checkpoints require trust_remote_code=True and ~48 GB GPU VRAM # (with 4-bit bitsandbytes quantization). Usage: - python turboquant/validate_minimax_m2.py - python turboquant/validate_minimax_m2.py --dry-run # synthetic benchmark only + python turboquant/validate_minimax_m2.py # M3 (default) + python turboquant/validate_minimax_m2.py --model MiniMax-M2.7 # legacy M2.7 + python turboquant/validate_minimax_m2.py --dry-run # synthetic only """ import argparse @@ -33,8 +35,9 @@ from turboquant.planarquant import PlanarQuantMSE from turboquant.compressors import TurboQuantCompressorV2, TurboQuantCompressorMSE -# ── MiniMax-M2.7 architecture constants ───────────────────────────── -MODEL_NAME = "MiniMaxAI/MiniMax-M2.7" +# ── MiniMax architecture constants (M3 / M2.7 share the same MoE layout) ─── +SUPPORTED_MODELS = ("MiniMax-M3", "MiniMax-M2.7") +DEFAULT_MODEL = "MiniMax-M3" HEAD_DIM = 128 NUM_KV_HEADS = 8 NUM_ATTN_HEADS = 48 @@ -122,7 +125,7 @@ def _compress_and_score_turbo(keys: torch.Tensor, query: torch.Tensor, def run_synthetic_benchmark(seq_len: int = 2048, n_layers: int = 8): """ - Benchmark on synthetic KV tensors with MiniMax-M2.7 dimensions. + Benchmark on synthetic KV tensors with MiniMax M3 / M2.7 dimensions. No model loading required — useful for CI / quick sanity checks. """ @@ -168,14 +171,15 @@ def run_synthetic_benchmark(seq_len: int = 2048, n_layers: int = 8): f"top1: {100*turbo_top1/turbo_n:.1f}%") -# ── Full validation with real MiniMax-M2.7 model ────────────────────── +# ── Full validation with real MiniMax model ─────────────────────────── -def run_model_validation(target_tokens: int = 2048): +def run_model_validation(model_id: str = DEFAULT_MODEL, target_tokens: int = 2048): """ - Run KV cache compression validation on a real MiniMax-M2.7 forward pass. + Run KV cache compression validation on a real MiniMax forward pass. Loads the model with 4-bit quantization (bitsandbytes NF4) to fit in - available GPU VRAM. + available GPU VRAM. Defaults to ``MiniMax-M3``; pass ``MiniMax-M2.7`` + via ``--model`` to validate against the previous flagship. """ try: from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig @@ -184,10 +188,11 @@ def run_model_validation(target_tokens: int = 2048): sys.exit(1) if not torch.cuda.is_available(): - print("ERROR: CUDA GPU required for MiniMax-M2.7 model validation.") + print(f"ERROR: CUDA GPU required for {model_id} model validation.") sys.exit(1) - print(f"\nLoading {MODEL_NAME} with 4-bit NF4 quantization…") + repo_id = f"MiniMaxAI/{model_id}" + print(f"\nLoading {repo_id} with 4-bit NF4 quantization…") bnb_cfg = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16, @@ -195,9 +200,9 @@ def run_model_validation(target_tokens: int = 2048): bnb_4bit_use_double_quant=True, ) - tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True) + tokenizer = AutoTokenizer.from_pretrained(repo_id, trust_remote_code=True) model = AutoModelForCausalLM.from_pretrained( - MODEL_NAME, + repo_id, quantization_config=bnb_cfg, device_map="auto", trust_remote_code=True, @@ -266,10 +271,10 @@ def run_model_validation(target_tokens: int = 2048): # ── Memory estimate ─────────────────────────────────────────────────── -def print_memory_estimate(seq_len: int = 32768): - """Print KV cache memory breakdown for MiniMax-M2.7 at various compressions.""" +def print_memory_estimate(seq_len: int = 32768, model_id: str = DEFAULT_MODEL): + """Print KV cache memory breakdown for MiniMax at various compressions.""" print(f"\n{'─' * 60}") - print(f"KV Cache Memory: MiniMax-M2.7, seq={seq_len:,}") + print(f"KV Cache Memory: {model_id}, seq={seq_len:,}") print(f" (kv_heads={NUM_KV_HEADS}, head_dim={HEAD_DIM}, layers={NUM_LAYERS})") print(f"{'─' * 60}") @@ -290,7 +295,11 @@ def print_memory_estimate(seq_len: int = 32768): def main(): parser = argparse.ArgumentParser( - description="Validate IsoQuant KV cache compression on MiniMax-M2.7" + description="Validate IsoQuant KV cache compression on MiniMax (M3 default)" + ) + parser.add_argument( + "--model", choices=SUPPORTED_MODELS, default=DEFAULT_MODEL, + help=f"MiniMax checkpoint to validate (default: {DEFAULT_MODEL})" ) parser.add_argument( "--dry-run", action="store_true", @@ -307,17 +316,17 @@ def main(): args = parser.parse_args() print("=" * 60) - print("RotorQuant × MiniMax-M2.7 KV Cache Compression") + print(f"RotorQuant × {args.model} KV Cache Compression") print("=" * 60) print(f"Model architecture: {HEAD_DIM}D head, {NUM_KV_HEADS} KV heads, " f"{NUM_LAYERS} layers, GQA") - print_memory_estimate(seq_len=32768) + print_memory_estimate(seq_len=32768, model_id=args.model) run_synthetic_benchmark(seq_len=2048, n_layers=args.layers) if not args.dry_run: - run_model_validation(target_tokens=args.seq_len) + run_model_validation(model_id=args.model, target_tokens=args.seq_len) print("\n" + "=" * 60) print("DONE")