chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 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

4
docs/supported_models/generative_models.md
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. |
@@ -64,4 +64,4 @@ in the GitHub search bar.
| **NVIDIA Nemotron Nano 2.0** | `nvidia/NVIDIA-Nemotron-Nano-9B-v2` | The [NVIDIA Nemotron](https://www.nvidia.com/en-us/ai-data-science/foundation-models/nemotron/) family of multimodal models provides state-of-the-art reasoning models specifically designed for enterprise-ready AI agents. `Nemotron-Nano-9B-v2` is a hybrid Mamba-Transformer language model designed to increase throughput for reasoning workloads while achieving state-of-the-art accuracy compared to similarly-sized models. |
| **StarCoder2** (3B-15B) | `bigcode/starcoder2-7b` | StarCoder2 is a family of open large language models (LLMs) specialized for code generation and understanding. It is the successor to StarCoder, jointly developed by the BigCode project (a collaboration between Hugging Face, ServiceNow Research, and other contributors). |
| **Jet-Nemotron** | `jet-ai/Jet-Nemotron-2B` | Jet-Nemotron is a new family of hybrid-architecture language models that surpass state-of-the-art open-source full-attention language models, while achieving significant efficiency gains. |
| **Trinity** (Nano, Mini) | `arcee-ai/Trinity-Mini` | Arcee's foundational MoE Trinity family of models, open weights under Apache 2.0. |
| **Trinity** (Nano, Mini) | `arcee-ai/Trinity-Mini` | Arcee's foundational MoE Trinity family of models, open weights under Apache 2.0. |

20
python/sglang/srt/models/commandr.py
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,7 +246,9 @@ 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)
q, k = self.rotary_emb(positions, 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)
return 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)
)