PromptCoT-Mamba-Math-7B

Scaling Reasoning without Attention

---

🚀 Overview

PromptCoT-Mamba establishes the first attention-free foundation model capable of surpassing strong Transformer baselines across a broad suite of competition-level math and code reasoning tasks. Built on the Mamba-2 architecture and trained through a structured, two-stage curriculum using the PromptCoT pipeline, it delivers high accuracy with constant-memory inference, eliminating the need for KV caching.

---

📈 Key Results

🔹 General Performance

🔍 PromptCoT-Mamba-7B consistently outperforms all 7B-scale Transformer and hybrid Mamba-Transformer baselines across all tasks.

---

🔹 Math Specialization vs. Generalist

🎯 The math-specialized variant improves AIME 24 by +7.7% and AIME 25 by +6.2%, with a slight trade-off in code-related performance.

---

⚡ Inference Efficiency

Using vLLM under constrained memory, PromptCoT-Mamba-7B demonstrates substantial speedups over the S1.1-7B Transformer baseline:

• 💡 3.66× faster at long-sequence generation on 24GB GPU
• 💡 1.69× faster under 72GB memory

⚙️ Practical for cost-sensitive or long-context inference workloads at scale.

---

🧪 Quick Start

🔧 Install Requirements

pip install transformers vllm torch accelerate

🧠 Load and Run the Model

from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "xl-zhao/PromptCoT-Mamba-Math-7B"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name).to("cuda")

problem_statement = (
    "A robe takes 2 bolts of blue fiber and half that much white fiber.  How many bolts in total does it take?"
)

prompt = (
    f"<|im_start|>user\n{problem_statement}\nPlease reason step by step, and put your final answer within \\boxed{{}}.<|im_end|>\n"
    "<|im_start|>assistant\n"
)

inputs = tokenizer(prompt, return_tensors="pt").to("cuda")

with torch.no_grad():
    output = model.generate(**inputs, max_length=65536, temperature=0.8)

generated_solution = tokenizer.decode(output[0], skip_special_tokens=True)
print(generated_solution)

---

⚡ Fast Inference with vLLM

from vllm import LLM, SamplingParams

model_name = "xl-zhao/PromptCoT-Mamba-Math-7B"
llm = LLM(model=model_name, tensor_parallel_size=1)

problem_statement = (
    "A robe takes 2 bolts of blue fiber and half that much white fiber.  How many bolts in total does it take?"
)

prompt = (
    f"<|im_start|>user\n{problem_statement}\nPlease reason step by step, and put your final answer within \\boxed{{}}.<|im_end|>\n"
    "<|im_start|>assistant\n"
)

sampling_params = SamplingParams(temperature=0.8, max_tokens=65536)
outputs = llm.generate([prompt], sampling_params)

print(outputs[0].outputs[0].text)

---

📜 Citation

@article{zhao2025scaling,
  author    = {Xueliang Zhao and Wei Wu and Lingpeng Kong},
  title     = {Scaling Reasoning without Attention},
  journal   = {arXiv preprint arXiv:2505.22425},
  year      = {2025},
  url       = {https://arxiv.org/abs/2505.22425}
}