minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[new-model] Add support for Cohere2ForCausalLM behind Command-A and Command-R Models (#16927)

Browse Source
This commit is contained in:
Lingjun Wen
2026-01-21 12:28:33 -08:00
committed by GitHub
parent 1fdf5cac39
commit cf89351691
2 changed files with 19 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/supported_models/generative_models.md
+1 -1
View File
@@ -39,7 +39,7 @@ in the GitHub search bar.
| **OLMoE** (Open MoE) | `allenai/OLMoE-1B-7B-0924` | Allen AI’s open Mixture-of-Experts model (7B total, 1B active parameters) delivering state-of-the-art results with sparse expert activation. |
| **MiniMax-M2** (M2, M2.1) | `minimax/MiniMax-M2`, `minimax/MiniMax-M2.1` | MiniMax’s SOTA LLM for coding & agentic workflows. |
| **StableLM** (3B, 7B) | `stabilityai/stablelm-tuned-alpha-7b` | StabilityAI’s early open-source LLM (3B & 7B) for general text generation; a demonstration model with basic instruction-following ability. |
| **Command-R** (Cohere) | `CohereForAI/c4ai-command-r-v01` | Cohere’s open conversational LLM (Command series) optimized for long context, retrieval-augmented generation, and tool use. |
| **Command-(R,A)** (Cohere) | `CohereLabs/c4ai-command-r-v01`, `CohereLabs/c4ai-command-r7b-12-2024`, `CohereLabs/c4ai-command-a-03-2025` | Cohere’s open conversational LLM (Command series) optimized for long context, retrieval-augmented generation, and tool use. |
| **DBRX** (Databricks) | `databricks/dbrx-instruct` | Databricks’ 132B-parameter MoE model (36B active) trained on 12T tokens; competes with GPT-3.5 quality as a fully open foundation model. |
| **Grok** (xAI) | `xai-org/grok-1` | xAI’s grok-1 model known for vast size(314B parameters) and high quality; integrated in SGLang for high-performance inference. |
| **ChatGLM** (GLM-130B family) | `THUDM/chatglm2-6b` | Zhipu AI’s bilingual chat model (6B) excelling at Chinese-English dialogue; fine-tuned for conversational quality and alignment. |
python/sglang/srt/models/commandr.py
+16 -2
View File
@@ -43,7 +43,7 @@ import torch
import torch.utils.checkpoint
from torch import nn
from torch.nn.parameter import Parameter
from transformers import PretrainedConfig
from transformers import Cohere2Config, CohereConfig, PretrainedConfig
from sglang.srt.distributed import (
get_tensor_model_parallel_rank,
@@ -198,12 +198,23 @@ class CohereAttention(nn.Module):
rope_scaling=self.rope_scaling,
is_neox_style=False,
)
self.v1 = isinstance(config, CohereConfig)
self.v2 = isinstance(config, Cohere2Config)
# Model v2 has interleaved sliding windows, v1 does not
if self.v2 and config.layer_types[layer_id] == "sliding_attention":
self.sliding_window_size = config.sliding_window
else:
self.sliding_window_size = -1
self.attn = RadixAttention(
self.num_heads,
self.head_dim,
self.scaling,
num_kv_heads=self.num_kv_heads,
layer_id=layer_id,
sliding_window_size=self.sliding_window_size,
quant_config=quant_config,
prefix=add_prefix("attn", prefix),
)
@@ -235,6 +246,8 @@ class CohereAttention(nn.Module):
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
if self.use_qk_norm:
q, k = self._apply_qk_norm(q, k)
# Model v1 uses RoPE throughout, Model v2 uses RoPE only for SWA layers
if self.v1 or self.sliding_window_size > 0:
q, k = self.rotary_emb(positions, q, k)
attn_output = self.attn(q, k, v, forward_batch)
output, _ = self.o_proj(attn_output)
@@ -348,7 +361,8 @@ class CohereForCausalLM(nn.Module):
super().__init__()
self.config = config
self.quant_config = quant_config
self.logits_processor = LogitsProcessor(config)
self.logit_scale = getattr(config, "logit_scale", None)
self.logits_processor = LogitsProcessor(config, logit_scale=self.logit_scale)
self.model = CohereModel(
config, quant_config, prefix=add_prefix("model", prefix)
)