Linked issue

Fixes #6684

Type of change

• New feature or public API

Motivation

GritLM-7B is a Mistral-7B-based model trained with Generative Representational
Instruction Tuning (GRIT), enabling both high-quality text generation and dense
vector embeddings from a single model. It is widely used in retrieval-augmented
generation (RAG) pipelines and semantic search applications.

MAX had no native support for the GritLM architecture class. This PR adds a implementation so parasail-ai/GritLM-7B-vllm can be served directly via max serve without any
custom flags after registration.

What changed

Added max/python/max/pipelines/architectures/gritlm/ — a new ModuleV3
architecture package for the GritLM model family.

New files:

Architecture highlights:

• Mistral-7B backbone: 32 layers, hidden_size=4096, GQA (32 Q / 8 KV heads),
  SwiGLU MLP, RMSNorm.
• Sliding window attention on all 32 layers (sliding_window=4096) using
  flash_attention_ragged with MHAMaskVariant.SLIDING_WINDOW_CAUSAL_MASK.
• Standard Mistral RoPE (rope_theta=10000, no scaling).
• Separate lm_head (tie_word_embeddings=false).
• CausalLM path only — gritlm_pooling_head.weight is dropped by the
  weight adapter since MAX serves text generation alone for this model.

Testing

Verified on parasail-ai/GritLM-7B-vllm with GPU serving:

max serve \
  --model-path parasail-ai/GritLM-7B-vllm \
  --custom-architectures architectures/gritlm \
  --max-batch-size 256 \
  --max-length 4096 \
  --quantization-encoding bfloat16

Smoke test — model generates correctly :

curl http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{"model":"parasail-ai/GritLM-7B-vllm",
       "messages":[{"role":"user","content":"What is 2+2?"}],
       "max_tokens":32,"temperature":0}'

Output :

2+2 is equal to 4.

GSM8K accuracy vs vLLM reference :

Checklist

• The linked issue above has been reviewed by a maintainer and is agreed-upon,
  or this is a trivial fix that does not need prior approval
• PR is small and focused — single new architecture, no changes to existing
  architectures
• I ran ./bazelw run format to format my changes
• I added or updated tests to cover my changes
• If AI tools assisted with this contribution, I have included an
  Assisted-by: trailer in my commit message or this PR description

Assisted-by: AI