diff --git a/docs/cookbook/autoregressive/RedNote/Dots3-Note.mdx b/docs/cookbook/autoregressive/RedNote/Dots3-Note.mdx index feb71c55a..dd6e990e2 100644 --- a/docs/cookbook/autoregressive/RedNote/Dots3-Note.mdx +++ b/docs/cookbook/autoregressive/RedNote/Dots3-Note.mdx @@ -57,6 +57,14 @@ import { config } from "/src/snippets/configs/rednote/dots3-note.jsx"; +## Playground + +The Playground is where you experiment with **SGLang features beyond the verified matrix**. The Deploy panel above only emits combinations signed off on this page; the Playground lets you turn on additional knobs on top of whichever cell the Deploy panel is currently showing. + +import { Playground } from "/src/snippets/_playground.jsx"; + + + ## 1. Model Introduction dots3.note is RedNote's native multimodal omni model, built on the dots3 language model. It accepts text, image, audio, and native video input. @@ -66,7 +74,12 @@ dots3.note is RedNote's native multimodal omni model, built on the dots3 languag - **Hybrid attention** — dots3 combines MLA with full-attention and sliding-window layers of different geometry, attention gates, and optional DSA indexing on full-attention layers. - **MTP speculative decoding** — a full-sharing MTP/NextN architecture exposes one recursively shared, SWA-shaped MTP layer and shares the target LM head. -**Resources:** [Hugging Face](https://huggingface.co/dots-studio/dots3-note-prev) · [SGLang PR #33829](https://github.com/sgl-project/sglang/pull/33829) +**Available checkpoints:** + +- **BF16**: [dots-studio/dots3-note-prev](https://huggingface.co/dots-studio/dots3-note-prev) +- **FP8**: [dots-studio/dots3-note-prev-fp8](https://huggingface.co/dots-studio/dots3-note-prev-fp8) + +**Resources:** [Hugging Face (BF16)](https://huggingface.co/dots-studio/dots3-note-prev) · [Hugging Face (FP8)](https://huggingface.co/dots-studio/dots3-note-prev-fp8) · [SGLang PR #33829](https://github.com/sgl-project/sglang/pull/33829) ## 2. Configuration Tips @@ -112,10 +125,12 @@ response = client.chat.completions.create( } ], extra_body={ - "seq": 131072, - "audio_cap": 0.5, - "audio_sr": 16000, - "k_mode": "eval_ek", + "video_config": { + "seq": 131072, + "audio_cap": 0.5, + "audio_sr": 16000, + "k_mode": "eval_ek", + }, }, ) @@ -132,28 +147,93 @@ Pending update... -Per-request video preprocessing controls (all optional, passed via `extra_body`): +Per-request video preprocessing controls are grouped under `video_config` in +`extra_body`: -| Field | Default | Purpose | -|-------|---------|---------| -| `seq` | `131072` | Total sequence budget used by the video flattener. | -| `audio_cap` | `1.0` | Maximum fraction of the input budget assigned to audio; `0` disables audio processing. | -| `audio_sr` | `16000` | Audio sample rate. | -| `k_mode` | `eval_ek` | Deterministic evaluation/sampling mode of the flattener. | + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
FieldDefaultPurpose
seq131072Total sequence budget used by the video flattener.
audio_cap1.0Maximum fraction of the input budget assigned to audio; 0 disables audio processing.
audio_sr16000Audio sample rate.
k_modeeval_ekDeterministic evaluation/sampling mode of the flattener.
-These controls are request-scoped so that evaluation jobs with different context budgets can share one server. The flattener reserves room for `max_new_tokens` inside the budget and falls back to visual-only processing if audio would exceed the configured token budget. +These controls are request-scoped so that evaluation jobs with different context budgets can share one server. For example: `extra_body={"video_config": {"seq": 131072, "audio_cap": 0.5}}`. The flattener reserves room for `max_new_tokens` inside the budget and falls back to visual-only processing if audio would exceed the configured token budget. - -Native video currently supports one video per request, and a native video cannot be mixed with separate image or audio inputs in the same request. - +A request may carry several videos, and videos can be mixed with image and audio parts. Each video is flattened independently under the same per-request budget, and the flattened frames and audio segments are spliced back at the position of their `video_url` part, so the modality ordering of the prompt is preserved. ### 3.2 Image and audio input -Outside the native-video path, images and audio clips use the standard OpenAI multimodal message format and SGLang's multimodal serving (`--enable-multimodal` is in every cell). The vision and audio towers run in-process, so no extra server is needed. +Outside the native-video path, images and audio clips use the standard OpenAI multimodal message format. `--enable-multimodal` is in every cell; the vision and audio towers run in-process, so no extra server is needed. + + + +```python Example +from openai import OpenAI + +client = OpenAI(base_url="http://localhost:30000/v1", api_key="EMPTY") + +response = client.chat.completions.create( + model="dots3.note", + messages=[ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": {"url": "https://example.com/sample.jpg"}, + }, + {"type": "text", "text": "Describe this image."}, + ], + } + ], +) + +print(response.choices[0].message.content) +``` + + + + + +```text Output +Pending update... +``` + + ### 3.3 Tool Calling -The cells launch with `--tool-call-parser dots`, so structured tool calls surface via `message.tool_calls` out of the box. +Toggle **Tool Call Parser** (`--tool-call-parser dots`) and **Reasoning Parser** (`--reasoning-parser dots`) in the **Parsers** card of the [Playground above](#playground). Structured tool calls then surface via `message.tool_calls`. @@ -195,7 +275,7 @@ Pending update... ### 3.4 Encoder/LLM Disaggregation (EPD) -`Dot3NoteForCausalLM` supports both roles of an encoder/LLM-disaggregated deployment: +`Dots3NoteForCausalLM` supports both roles of an encoder/LLM-disaggregated deployment: - **Encoder role** — serve with `--encoder-only`; the instance runs only the vision and audio towers. - **Language role** — serve with `--language-only`; the instance skips tower construction, leaving the memory to the language model. diff --git a/docs/src/snippets/configs/rednote/dots3-note.jsx b/docs/src/snippets/configs/rednote/dots3-note.jsx index 33ed19a8e..8e85f19cd 100644 --- a/docs/src/snippets/configs/rednote/dots3-note.jsx +++ b/docs/src/snippets/configs/rednote/dots3-note.jsx @@ -1,13 +1,9 @@ // Dots3-Note cookbook config. Consumed by _deployment.jsx + _playground.jsx. -// Single `export const config` literal - no spreads/calls/IIFE (Mintlify re-evals at hydration). +// Single `export const config` literal — no spreads/calls/IIFE (Mintlify re-evals at hydration). export const config = { modelName: "Dots3-Note", - // No Playground on this page — the only extra knob (the dots tool-call parser) - // is already baked into the cells. - showPlaygroundLink: false, - // Hopper only for now — no Blackwell support. supportedHardware: ["h200", "h100"], @@ -50,14 +46,78 @@ export const config = { {"type": "video_url", "video_url": {"url": "https://example.com/sample.mp4"}}, {"type": "text", "text": "Summarize what happens in this video."} ] - }] + }], + "video_config": { + "seq": 131072, + "audio_cap": 0.5, + "audio_sr": 16000, + "k_mode": "eval_ek" + } }'`, dockerImages: { - h200: "lmsysorg/sglang:dev", - h100: "lmsysorg/sglang:dev", + h200: "lmsysorg/sglang:dev-dots3-note", + h100: "lmsysorg/sglang:dev-dots3-note", }, + github: { + cookbookModel: "dots-studio/dots3-note-prev", + }, + + playgroundFeatures: { + attention: { + knobs: [ + { id: "tp", label: "TP", values: [null, 4, 8] }, + { + id: "dpAttn", + label: "DP-Attention", + values: [null, false, 4, 8], + labels: { "auto": "Auto", "false": "Off" }, + }, + ], + }, + moe: { + backend: { + options: [ + { id: null, label: "Inherited" }, + { id: "deepep", label: "DeepEP", flags: ["--moe-a2a-backend deepep"] }, + ], + }, + ep: { label: "EP", values: [null, 4, 8] }, + }, + parsers: { + items: [ + { + id: "reasoning", + label: "Reasoning Parser", + flag: "--reasoning-parser dots", + }, + { + id: "toolCall", + label: "Tool Call Parser", + flag: "--tool-call-parser dots", + }, + ], + }, + speculative: { + options: [ + { id: "current", label: "Inherited from base" }, + { id: "off", label: "Off (greedy)" }, + { + id: "nextn-314", + label: "NEXTN / MTP 3-1-4", + flags: [ + "--speculative-algorithm NEXTN", + "--speculative-num-steps 3", + "--speculative-eagle-topk 1", + "--speculative-num-draft-tokens 4", + "--speculative-draft-model-path {{MODEL_NAME}}", + "--speculative-draft-attention-backend fa3", + ], + }, + ], + }, + }, cells: [ { @@ -74,20 +134,25 @@ export const config = { ], flags: [ "--model-path {{MODEL_NAME}}", - "--context-length 524288", + "--trust-remote-code", "--enable-dp-attention", - "--dp-size 8", - "--tp-size 8", - "--ep-size 8", + "--tp 8", + "--dp 8", + "--ep 8", + "--moe-dense-tp-size 1", + "--moe-a2a-backend deepep", + "--moe-runner-backend deep_gemm", + "--deepep-dispatcher-output-dtype bf16", + "--deepep-mode auto", + "--enable-nccl-nvls", + "--context-length 524288", "--mem-fraction-static 0.87", "--max-running-requests 256", "--chunked-prefill-size 16384", - "--trust-remote-code", "--swa-full-tokens-ratio 0.03", "--prefill-attention-backend fa3", "--decode-attention-backend fa3", "--page-size 64", - "--moe-dense-tp-size 1", "--cuda-graph-backend-decode full", "--cuda-graph-backend-prefill disabled", "--cuda-graph-max-bs-decode 32", @@ -97,15 +162,8 @@ export const config = { "--speculative-num-draft-tokens 4", "--speculative-draft-model-path {{MODEL_NAME}}", "--speculative-draft-attention-backend fa3", - "--moe-a2a-backend deepep", - "--moe-runner-backend deep_gemm", - "--deepep-dispatcher-output-dtype bf16", - "--deepep-mode auto", - "--enable-nccl-nvls", "--enable-multimodal", "--enable-metrics", - "--tool-call-parser dots", - "--reasoning-parser qwen3", "--watchdog-timeout 1800", "--host {{HOST_IP}}", "--port {{PORT}}", @@ -125,20 +183,25 @@ export const config = { ], flags: [ "--model-path {{MODEL_NAME}}", - "--context-length 524288", + "--trust-remote-code", "--enable-dp-attention", - "--dp-size 8", - "--tp-size 8", - "--ep-size 8", + "--tp 8", + "--dp 8", + "--ep 8", + "--moe-dense-tp-size 1", + "--moe-a2a-backend deepep", + "--moe-runner-backend auto", + "--deepep-dispatcher-output-dtype auto", + "--deepep-mode auto", + "--enable-nccl-nvls", + "--context-length 524288", "--mem-fraction-static 0.87", "--max-running-requests 256", "--chunked-prefill-size 16384", - "--trust-remote-code", "--swa-full-tokens-ratio 0.03", "--prefill-attention-backend fa3", "--decode-attention-backend fa3", "--page-size 64", - "--moe-dense-tp-size 1", "--cuda-graph-backend-decode full", "--cuda-graph-backend-prefill disabled", "--cuda-graph-max-bs-decode 32", @@ -148,15 +211,8 @@ export const config = { "--speculative-num-draft-tokens 4", "--speculative-draft-model-path {{MODEL_NAME}}", "--speculative-draft-attention-backend fa3", - "--moe-a2a-backend deepep", - "--moe-runner-backend auto", - "--deepep-dispatcher-output-dtype auto", - "--deepep-mode auto", - "--enable-nccl-nvls", "--enable-multimodal", "--enable-metrics", - "--reasoning-parser qwen3", - "--tool-call-parser dots", "--watchdog-timeout 1800", "--host {{HOST_IP}}", "--port {{PORT}}", @@ -176,20 +232,25 @@ export const config = { ], flags: [ "--model-path {{MODEL_NAME}}", - "--context-length 524288", + "--trust-remote-code", "--enable-dp-attention", - "--dp-size 8", - "--tp-size 8", - "--ep-size 8", + "--tp 8", + "--dp 8", + "--ep 8", + "--moe-dense-tp-size 1", + "--moe-a2a-backend deepep", + "--moe-runner-backend deep_gemm", + "--deepep-dispatcher-output-dtype bf16", + "--deepep-mode auto", + "--enable-nccl-nvls", + "--context-length 524288", "--mem-fraction-static 0.87", "--max-running-requests 256", "--chunked-prefill-size 16384", - "--trust-remote-code", "--swa-full-tokens-ratio 0.03", "--prefill-attention-backend fa3", "--decode-attention-backend fa3", "--page-size 64", - "--moe-dense-tp-size 1", "--cuda-graph-backend-decode full", "--cuda-graph-backend-prefill disabled", "--cuda-graph-max-bs-decode 32", @@ -199,15 +260,8 @@ export const config = { "--speculative-num-draft-tokens 4", "--speculative-draft-model-path {{MODEL_NAME}}", "--speculative-draft-attention-backend fa3", - "--moe-a2a-backend deepep", - "--moe-runner-backend deep_gemm", - "--deepep-dispatcher-output-dtype bf16", - "--deepep-mode auto", - "--enable-nccl-nvls", "--enable-multimodal", "--enable-metrics", - "--tool-call-parser dots", - "--reasoning-parser qwen3", "--watchdog-timeout 1800", "--host {{HOST_IP}}", "--port {{PORT}}", @@ -227,20 +281,25 @@ export const config = { ], flags: [ "--model-path {{MODEL_NAME}}", - "--context-length 524288", + "--trust-remote-code", "--enable-dp-attention", - "--dp-size 8", - "--tp-size 8", - "--ep-size 8", + "--tp 8", + "--dp 8", + "--ep 8", + "--moe-dense-tp-size 1", + "--moe-a2a-backend deepep", + "--moe-runner-backend auto", + "--deepep-dispatcher-output-dtype auto", + "--deepep-mode auto", + "--enable-nccl-nvls", + "--context-length 524288", "--mem-fraction-static 0.87", "--max-running-requests 256", "--chunked-prefill-size 16384", - "--trust-remote-code", "--swa-full-tokens-ratio 0.03", "--prefill-attention-backend fa3", "--decode-attention-backend fa3", "--page-size 64", - "--moe-dense-tp-size 1", "--cuda-graph-backend-decode full", "--cuda-graph-backend-prefill disabled", "--cuda-graph-max-bs-decode 32", @@ -250,15 +309,8 @@ export const config = { "--speculative-num-draft-tokens 4", "--speculative-draft-model-path {{MODEL_NAME}}", "--speculative-draft-attention-backend fa3", - "--moe-a2a-backend deepep", - "--moe-runner-backend auto", - "--deepep-dispatcher-output-dtype auto", - "--deepep-mode auto", - "--enable-nccl-nvls", "--enable-multimodal", "--enable-metrics", - "--reasoning-parser qwen3", - "--tool-call-parser dots", "--watchdog-timeout 1800", "--host {{HOST_IP}}", "--port {{PORT}}", diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index 8722782f3..167959763 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -599,6 +599,7 @@ def _kimi_k3_moe_runner_overrides(server_args: Any, hf_config: Any) -> dict: "GlmMoeDsaForCausalLM", "LongcatFlashForCausalLM", "LongcatFlashForCausalLMNextN", + "Dots3NoteForCausalLM", ) def _deepseek_family_overrides(server_args: Any, hf_config: Any) -> dict: """Order-safe declarations of the DeepSeek/DSA branch. The CP parallel @@ -609,6 +610,7 @@ def _deepseek_family_overrides(server_args: Any, hf_config: Any) -> dict: from sglang.srt.configs.model_config import is_deepseek_dsa overrides: Dict[str, Any] = {} + if is_deepseek_dsa(hf_config): # DeepSeek 3.2/GLM 5 # Set attention backend for DeepSeek if server_args.is_attention_backend_not_set(): @@ -1769,6 +1771,7 @@ _DEEPSEEK_FAMILY_ARCHS = frozenset( "GlmMoeDsaForCausalLM", "LongcatFlashForCausalLM", "LongcatFlashForCausalLMNextN", + "Dots3NoteForCausalLM", } ) diff --git a/python/sglang/srt/configs/__init__.py b/python/sglang/srt/configs/__init__.py index 54b5fa46c..3e9b06ce2 100644 --- a/python/sglang/srt/configs/__init__.py +++ b/python/sglang/srt/configs/__init__.py @@ -4,6 +4,7 @@ from sglang.srt.configs.chatglm import ChatGLMConfig from sglang.srt.configs.cohere2_moe import Cohere2MoeConfig from sglang.srt.configs.dbrx import DbrxConfig from sglang.srt.configs.deepseekvl2 import DeepseekVL2Config +from sglang.srt.configs.dots3 import Dots3Config from sglang.srt.configs.dots_ocr import DotsOCRConfig from sglang.srt.configs.dots_vlm import DotsVLMConfig from sglang.srt.configs.exaone import ExaoneConfig @@ -97,6 +98,7 @@ __all__ = [ "InternS2MobiusVisionConfig", "DotsVLMConfig", "DotsOCRConfig", + "Dots3Config", "FalconH1Config", "GraniteMoeHybridConfig", "Lfm2Config", diff --git a/python/sglang/srt/configs/dots3.py b/python/sglang/srt/configs/dots3.py new file mode 100644 index 000000000..22766f4c2 --- /dev/null +++ b/python/sglang/srt/configs/dots3.py @@ -0,0 +1,243 @@ +from transformers import AutoTokenizer +from transformers.configuration_utils import PretrainedConfig + +from sglang.srt.multimodal.customized_mm_processor_utils import ( + register_customized_processor, +) + + +class DotsNoteOmniTokenizerProxy: + @classmethod + def from_pretrained(cls, model_path: str, *args, **kwargs): + kwargs.pop("use_fast", None) + return AutoTokenizer.from_pretrained(model_path, *args, **kwargs) + + +@register_customized_processor(DotsNoteOmniTokenizerProxy) +class Dots3Config(PretrainedConfig): + model_type = "dots3_note" + keys_to_ignore_at_inference = ["past_key_values"] + is_hybrid_swa = True + requires_draft_attention_wrapper = True + + def __init__( + self, + # General model parameters + vocab_size=152064, + hidden_size=2560, + hidden_act="silu", + intermediate_size=7168, + num_hidden_layers=30, + max_position_embeddings=8192, + initializer_range=0.02, + rms_norm_eps=1e-5, + use_cache=True, + pretraining_tp=1, + # Token IDs + pad_token_id=None, + bos_token_id=151643, + eos_token_id=151645, + tie_word_embeddings=False, + # Attention parameters + attention_bias=False, + attention_dropout=0.0, + apply_mla_qkv_lora_rescale=True, + # MLA (Multi-head Latent Attention) parameters + attention_gate_type="headwise", + kv_lora_rank=512, + q_lora_rank=512, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + num_attention_heads=64, + num_key_value_heads=64, + v_head_dim=128, + # Dots3 uses one shared MTP layer for NEXTN decoding. + num_nextn_predict_layers=1, + # Sliding Window Attention (SWA) parameters + layer_types=None, + sliding_window_size=512, + swa_attention_gate_type="headwise", + swa_q_lora_rank=512, + swa_kv_lora_rank=512, + swa_qk_nope_head_dim=128, + swa_qk_rope_head_dim=64, + swa_rope_theta=None, + swa_num_attention_heads=32, + swa_num_key_value_heads=32, + swa_v_head_dim=128, + # MoE (Mixture of Experts) parameters + moe_intermediate_size=1024, + n_shared_experts=1, + n_routed_experts=128, + num_experts_per_tok=6, + moe_layer_freq=1, + first_k_dense_replace=1, + routed_scaling_factor=1.0, + norm_topk_prob=True, + scoring_func="sigmoid", + n_group=1, + topk_method="noaux_tc", + topk_group=1, + # RoPE parameters + rope_theta=50000.0, + rope_scaling=None, + # Optional DSA indexer parameters. + index_n_heads=None, + index_head_dim=None, + index_topk=None, + language_only=False, + # Multimodal special tokens + im_start_token="<|img|>", + im_token="<|imgpad|>", + im_end_token="<|endofimg|>", + audio_start_token="<|audio_comp_start|>", + audio_token="<|audio_comp_pad|>", + audio_end_token="<|audio_comp_end|>", + video_token="<|video_pad|>", + **kwargs, + ): + # General model parameters + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.hidden_act = hidden_act + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.max_position_embeddings = max_position_embeddings + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.use_cache = use_cache + self.pretraining_tp = pretraining_tp + + # Attention parameters + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + self.apply_mla_qkv_lora_rescale = apply_mla_qkv_lora_rescale + + # MLA (Multi-head Latent Attention) parameters + self.attention_gate_type = attention_gate_type + self.kv_lora_rank = kv_lora_rank + self.q_lora_rank = q_lora_rank + self.qk_nope_head_dim = qk_nope_head_dim + self.qk_rope_head_dim = qk_rope_head_dim + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + self.v_head_dim = v_head_dim + + # MTP / NextN + self.num_nextn_predict_layers = num_nextn_predict_layers + + # Sliding Window Attention (SWA) parameters + self.layer_types = layer_types + self.sliding_window_size = sliding_window_size + self.swa_attention_gate_type = swa_attention_gate_type + self.swa_q_lora_rank = swa_q_lora_rank + self.swa_kv_lora_rank = swa_kv_lora_rank + self.swa_qk_nope_head_dim = swa_qk_nope_head_dim + self.swa_qk_rope_head_dim = swa_qk_rope_head_dim + self.swa_rope_theta = rope_theta if swa_rope_theta is None else swa_rope_theta + self.swa_num_attention_heads = swa_num_attention_heads + self.swa_num_key_value_heads = swa_num_key_value_heads + self.swa_v_head_dim = swa_v_head_dim + # Runtime cache geometry for the SWA attention path. + self.swa_head_dim = swa_qk_nope_head_dim + swa_qk_rope_head_dim + + # MoE (Mixture of Experts) parameters + self.moe_intermediate_size = moe_intermediate_size + self.n_shared_experts = n_shared_experts + self.n_routed_experts = n_routed_experts + self.num_experts_per_tok = num_experts_per_tok + self.moe_layer_freq = moe_layer_freq + self.first_k_dense_replace = first_k_dense_replace + self.routed_scaling_factor = routed_scaling_factor + self.norm_topk_prob = norm_topk_prob + self.scoring_func = scoring_func + self.n_group = n_group + self.topk_method = topk_method + self.topk_group = topk_group + + # RoPE parameters + self.rope_theta = rope_theta + self.rope_scaling = rope_scaling + self._rope_scaling_validation() + + # NSA (Native Sparse Attention) parameters + self.index_n_heads = index_n_heads + self.index_head_dim = index_head_dim + self.index_topk = index_topk + self.language_only = language_only + + self.im_start_token = im_start_token + self.im_token = im_token + self.im_end_token = im_end_token + self.audio_start_token = audio_start_token + self.audio_token = audio_token + self.audio_end_token = audio_end_token + # The chat template renders a video content part as this single token, + # which the processor replaces with the flattened frames and audio. + self.video_token = video_token + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) + + def configure_draft_model(self) -> str: + """Configure the recursively shared MTP layer with SWA geometry.""" + self.num_nextn_predict_layers = 1 + self.layer_types = ["sliding_attention"] + self.attention_gate_type = self.swa_attention_gate_type + self.kv_lora_rank = self.swa_kv_lora_rank + self.q_lora_rank = self.swa_q_lora_rank + self.qk_nope_head_dim = self.swa_qk_nope_head_dim + self.qk_rope_head_dim = self.swa_qk_rope_head_dim + self.num_attention_heads = self.swa_num_attention_heads + self.num_key_value_heads = self.swa_num_key_value_heads + self.v_head_dim = self.swa_v_head_dim + return "Dots3NoteForCausalLMNextN" + + def wrap_attention_backend(self, runner, full_attn_backend): + from sglang.srt.layers.attention.dots_hybrid_backend import ( + wrap_dots_attention_backend, + ) + + return wrap_dots_attention_backend(runner, full_attn_backend) + + def wrap_draft_decode_attention_backend(self, backend): + from sglang.srt.layers.attention.dots_hybrid_backend import ( + wrap_dots_draft_decode_backend, + ) + + return wrap_dots_draft_decode_backend(backend) + + def _rope_scaling_validation(self): + """ + Validate the `rope_scaling` configuration. + """ + if self.rope_scaling is None: + return + + if not isinstance(self.rope_scaling, dict): + raise ValueError( + f"`rope_scaling` must be a dictionary, got {self.rope_scaling}" + ) + rope_scaling_type = self.rope_scaling.get("type", None) + rope_scaling_factor = self.rope_scaling.get("factor", None) + if rope_scaling_type is None or rope_scaling_type not in [ + "linear", + "dynamic", + "yarn", + ]: + raise ValueError( + f"`rope_scaling`'s type field must be one of ['linear', 'dynamic', 'yarn'], got {rope_scaling_type}" + ) + if ( + rope_scaling_factor is None + or not isinstance(rope_scaling_factor, (int, float)) + or rope_scaling_factor <= 1.0 + ): + raise ValueError( + f"`rope_scaling`'s factor field must be a number > 1, got {rope_scaling_factor}" + ) diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index fa8da9951..cbd054056 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -124,6 +124,8 @@ def is_deepseek_dsa(config) -> bool: "GlmMoeDsaForCausalLMNextN", "LongcatFlashForCausalLM", "LongcatFlashForCausalLMNextN", + "Dots3NoteForCausalLM", + "Dots3NoteForCausalLMNextN", ) and _hf_attr(config, "index_topk") is not None ) @@ -630,6 +632,13 @@ class ModelConfig: def _config_draft_model(self): is_draft_model = self.is_draft_model + from sglang.srt.configs.dots3 import Dots3Config + + if is_draft_model and isinstance(self.hf_text_config, Dots3Config): + self.hf_config.architectures[0] = ( + self.hf_text_config.configure_draft_model() + ) + if is_draft_model and self.hf_config.architectures[0] in [ "DeepseekV3ForCausalLM", "DeepseekV32ForCausalLM", @@ -866,6 +875,8 @@ class ModelConfig: self.hf_config.context_len = self.context_len def _derive_model_shapes(self): + from sglang.srt.configs.dots3 import Dots3Config + # Unify the config keys for hf_text_config self.head_dim = getattr(self.hf_text_config, "head_dim", None) if self.head_dim is None: @@ -902,6 +913,8 @@ class ModelConfig: or "LongcatFlashForCausalLM" in self.hf_config.architectures or "LongcatFlashForCausalLMNextN" in self.hf_config.architectures or "DotsVLMForCausalLM" in self.hf_config.architectures + or "Dots3NoteForCausalLM" in self.hf_config.architectures + or "Dots3NoteForCausalLMNextN" in self.hf_config.architectures or "MistralLarge3ForCausalLM" in self.hf_config.architectures or ( "PixtralForConditionalGeneration" in self.hf_config.architectures @@ -917,6 +930,12 @@ class ModelConfig: self.qk_nope_head_dim = self.hf_text_config.qk_nope_head_dim self.qk_rope_head_dim = self.hf_text_config.qk_rope_head_dim self.v_head_dim = self.hf_text_config.v_head_dim + if isinstance(self.hf_text_config, Dots3Config): + self.swa_kv_lora_rank = self.hf_text_config.swa_kv_lora_rank + self.swa_qk_rope_head_dim = self.hf_text_config.swa_qk_rope_head_dim + else: + self.swa_kv_lora_rank = self.kv_lora_rank + self.swa_qk_rope_head_dim = self.qk_rope_head_dim self.index_head_dim = ( get_dsa_index_head_dim(self.hf_text_config) if is_deepseek_dsa(self.hf_text_config) @@ -1893,6 +1912,7 @@ multimodal_model_archs = [ "Step3VLForConditionalGeneration", "POINTSV15ChatModel", "DotsVLMForCausalLM", + "Dots3NoteForCausalLM", "DotsOCRForCausalLM", "Sarashina2VisionForCausalLM", "NVILAForConditionalGeneration", @@ -2118,7 +2138,10 @@ def get_hybrid_layer_ids( full_attention_layer_ids = [ i for i in range(num_hidden_layers) if (i + 1) % 4 == 0 ] - elif any(arch in SWA_SINK_ARCHS for arch in model_architectures): + elif any(arch in SWA_SINK_ARCHS for arch in model_architectures) or any( + arch in ("Dots3NoteForCausalLM", "Dots3NoteForCausalLMNextN") + for arch in model_architectures + ): layer_types = getattr(hf_text_config, "layer_types", []) swa_attention_layer_ids = [ i for i, x in enumerate(layer_types) if x == "sliding_attention" diff --git a/python/sglang/srt/entrypoints/openai/protocol.py b/python/sglang/srt/entrypoints/openai/protocol.py index 656fe73e8..109542667 100644 --- a/python/sglang/srt/entrypoints/openai/protocol.py +++ b/python/sglang/srt/entrypoints/openai/protocol.py @@ -915,6 +915,7 @@ class ChatCompletionRequest(BaseModel): use_audio_in_video: bool = False images_config: Optional[Dict] = None + video_config: Optional[Dict] = None # Custom logit processor for advanced sampling control custom_logit_processor: Optional[Union[List[Optional[str]], str]] = None diff --git a/python/sglang/srt/entrypoints/openai/serving_chat.py b/python/sglang/srt/entrypoints/openai/serving_chat.py index 7bdc08b28..5055f9bb3 100644 --- a/python/sglang/srt/entrypoints/openai/serving_chat.py +++ b/python/sglang/srt/entrypoints/openai/serving_chat.py @@ -38,7 +38,10 @@ from jsonschema import Draft202012Validator, SchemaError from sglang.srt.entrypoints.openai import chat_encoding, encoding_dsv4, encoding_dsv32 from sglang.srt.entrypoints.openai.protocol import ( + ChatCompletionMessageContentTextPart, + ChatCompletionMessageContentVideoPart, ChatCompletionMessageGenericParam, + ChatCompletionMessageUserParam, ChatCompletionRequest, ChatCompletionResponse, ChatCompletionResponseChoice, @@ -205,6 +208,38 @@ def neutralize_kimi_k3_image_placeholder_value(value: Any) -> Any: return value +def _extract_video_question(request: ChatCompletionRequest) -> Optional[str]: + """Return text paired with a video in the last user turn.""" + for message in reversed(request.messages or []): + if not isinstance(message, ChatCompletionMessageUserParam): + continue + content = message.content + if not isinstance(content, list): + continue + has_video = any( + isinstance(part, ChatCompletionMessageContentVideoPart) for part in content + ) + if not has_video: + continue + return "".join( + part.text + for part in content + if isinstance(part, ChatCompletionMessageContentTextPart) + ) + return None + + +def _build_video_config(request: ChatCompletionRequest) -> Optional[Dict[str, Any]]: + """Build request-scoped video processor config without model-specific fields.""" + config = dict(request.video_config or {}) + question = _extract_video_question(request) + if question is not None: + # Internal metadata derived from the message must not be overridden by + # a model-specific public processor option. + config["_question"] = question + return config or None + + class OpenAIServingChat(OpenAIServingBase): """Handler for /v1/chat/completions requests""" @@ -1045,6 +1080,7 @@ class OpenAIServingChat(OpenAIServingBase): custom_labels=custom_labels, custom_logit_processor=request.custom_logit_processor, images_config=getattr(request, "images_config", None), + video_config=_build_video_config(request), image_max_dynamic_patch=img_max_dynamic_patch, video_max_dynamic_patch=vid_max_dynamic_patch, max_dynamic_patch=getattr(request, "max_dynamic_patch", None), diff --git a/python/sglang/srt/function_call/dots_detector.py b/python/sglang/srt/function_call/dots_detector.py new file mode 100644 index 000000000..26edddf82 --- /dev/null +++ b/python/sglang/srt/function_call/dots_detector.py @@ -0,0 +1,353 @@ +import json +import logging +import re +from typing import Any + +try: + import json_repair +except ImportError: + json_repair = None + +from sglang.srt.entrypoints.openai.protocol import Tool +from sglang.srt.function_call.base_format_detector import BaseFormatDetector +from sglang.srt.function_call.core_types import ( + StreamingParseResult, + StructureInfo, + ToolCallItem, + _GetInfoFunc, +) +from sglang.srt.function_call.utils import _is_complete_json + +logger = logging.getLogger(__name__) + + +class DotsToolDetector(BaseFormatDetector): + """Detector for the dots function-call format. + + The canonical format contains one or more XML ``invoke`` elements inside a + ``dots_function_call`` block:: + + + + weather in Shanghai + + + + A JSON object with ``name`` and ``arguments`` is accepted as a fallback. + Multiple wrapper blocks and multiple invokes in one block are supported. + """ + + def __init__(self): + super().__init__() + self.bot_token = "" + self.eot_token = "" + self.func_call_regex = re.compile( + rf"{re.escape(self.bot_token)}\s*(.*?)\s*{re.escape(self.eot_token)}", + re.DOTALL, + ) + self.invoke_regex = re.compile( + r"[^>]+)>(?P.*?)", + re.DOTALL, + ) + self.parameter_regex = re.compile( + r"[^>]+)>(?P.*?)", + re.DOTALL, + ) + + @staticmethod + def _extract_name(value: str) -> str: + value = value.strip() + if len(value) >= 2 and value[0] == value[-1] and value[0] in {'"', "'"}: + return value[1:-1] + return value + + @staticmethod + def _load_json(value: str) -> Any: + try: + return json.loads(value) + except (json.JSONDecodeError, ValueError): + if json_repair is None: + raise + return json_repair.loads(value) + + @classmethod + def _convert_param_value(cls, value: str, param_type: Any) -> Any: + if value.lower() == "null": + return None + + if isinstance(param_type, list): + param_type = next((item for item in param_type if item != "null"), "string") + if not isinstance(param_type, str): + param_type = str(param_type) + param_type = param_type.lower() + + if param_type in {"string", "str", "text"}: + return value + if param_type in {"integer", "int"}: + try: + return int(value) + except (TypeError, ValueError): + return value + if param_type in {"number", "float"}: + try: + number = float(value) + return int(number) if number.is_integer() else number + except (TypeError, ValueError): + return value + if param_type in {"boolean", "bool"}: + return value.lower() in {"true", "1"} + + try: + return cls._load_json(value) + except (json.JSONDecodeError, ValueError, TypeError): + return value + + def _resolve_param_type( + self, schema: Any, defs: dict[str, Any], depth: int = 0 + ) -> Any | None: + """Resolve a parameter type through local refs and schema compositions.""" + if not isinstance(schema, dict) or depth > 10: + return None + if "type" in schema: + return schema["type"] + + ref = schema.get("$ref") + if isinstance(ref, str) and ref.startswith("#/$defs/"): + return self._resolve_param_type( + defs.get(ref.rsplit("/", 1)[-1]), defs, depth + 1 + ) + + for keyword in ("anyOf", "oneOf", "allOf"): + alternatives = schema.get(keyword) + if not isinstance(alternatives, list): + continue + for alternative in alternatives: + if isinstance(alternative, dict) and alternative.get("type") == "null": + continue + resolved = self._resolve_param_type(alternative, defs, depth + 1) + if resolved is not None: + return resolved + return None + + @staticmethod + def _tool_schema(name: str, tools: list[Tool]) -> tuple[dict, dict]: + for tool in tools: + if tool.function.name != name: + continue + schema = tool.function.parameters + if not isinstance(schema, dict): + break + properties = schema.get("properties", {}) + defs = schema.get("$defs", {}) + return ( + properties if isinstance(properties, dict) else {}, + defs if isinstance(defs, dict) else {}, + ) + return {}, {} + + def _parse_xml_invoke(self, match: re.Match, tools: list[Tool]) -> dict[str, Any]: + name = self._extract_name(match.group("name")) + properties, defs = self._tool_schema(name, tools) + arguments: dict[str, Any] = {} + + for parameter in self.parameter_regex.finditer(match.group("body")): + param_name = self._extract_name(parameter.group("name")) + value = parameter.group("value").strip() + param_type: Any = "string" + if param_name in properties: + param_type = ( + self._resolve_param_type(properties[param_name], defs) or "string" + ) + arguments[param_name] = self._convert_param_value(value, param_type) + + return {"name": name, "arguments": arguments} + + def _parse_block(self, content: str, tools: list[Tool]) -> list[dict[str, Any]]: + content = content.strip() + if content.startswith(" bool: + return self.bot_token in text + + def detect_and_parse(self, text: str, tools: list[Tool]) -> StreamingParseResult: + marker_index = text.find(self.bot_token) + if marker_index == -1: + return StreamingParseResult(normal_text=text) + + calls: list[ToolCallItem] = [] + for block in self.func_call_regex.finditer(text): + try: + for parsed in self._parse_block(block.group(1), tools): + calls.extend(self.parse_base_json(parsed, tools)) + except (json.JSONDecodeError, ValueError, TypeError) as exc: + logger.warning("Failed to parse dots tool call: %s", exc) + + return StreamingParseResult( + normal_text=text[:marker_index].strip(), calls=calls + ) + + def _append_stream_call( + self, parsed: dict[str, Any], item: ToolCallItem + ) -> ToolCallItem: + self.current_tool_id += 1 + arguments = parsed.get("arguments", parsed.get("parameters", {})) or {} + serialized = json.dumps(arguments, ensure_ascii=False) + self.prev_tool_call_arr.append( + {"name": parsed.get("name"), "arguments": arguments} + ) + self.streamed_args_for_tool.append(serialized) + item.tool_index = self.current_tool_id + item.parameters = serialized + return item + + def parse_streaming_increment( + self, new_text: str, tools: list[Tool] + ) -> StreamingParseResult: + """Buffer incomplete XML and emit every complete call in the new data.""" + self._buffer += new_text + normal_parts: list[str] = [] + calls: list[ToolCallItem] = [] + + while self._buffer: + marker_index = self._buffer.find(self.bot_token) + if marker_index == -1: + partial_len = self._ends_with_partial_token( + self._buffer, self.bot_token + ) + if partial_len: + normal_parts.append(self._buffer[:-partial_len]) + self._buffer = self._buffer[-partial_len:] + else: + normal_parts.append(self._buffer) + self._buffer = "" + normal_parts = [ + part.replace(self.eot_token, "") for part in normal_parts + ] + break + + if marker_index > 0: + normal_parts.append(self._buffer[:marker_index]) + self._buffer = self._buffer[marker_index:] + + end_index = self._buffer.find(self.eot_token, len(self.bot_token)) + if end_index == -1: + self._stream_complete_json_body(tools, calls) + break + + content = self._buffer[len(self.bot_token) : end_index] + self._buffer = self._buffer[end_index + len(self.eot_token) :] + try: + parsed_calls = self._parse_block(content, tools) + if not parsed_calls: + raise ValueError("dots tool-call block contains no invoke") + block_calls: list[ToolCallItem] = [] + for index, parsed in enumerate(parsed_calls): + validated = self.parse_base_json(parsed, tools) + if index == 0 and self.current_tool_name_sent and validated: + item = validated[0] + arguments = item.parameters or "" + streamed = self.streamed_args_for_tool[self.current_tool_id] + remaining = arguments.removeprefix(streamed) + if remaining: + block_calls.append( + ToolCallItem( + tool_index=self.current_tool_id, + name=None, + parameters=remaining, + ) + ) + self.prev_tool_call_arr[self.current_tool_id] = parsed + self.streamed_args_for_tool[self.current_tool_id] = arguments + else: + block_calls.extend( + self._append_stream_call(parsed, item) for item in validated + ) + if block_calls: + calls.extend(block_calls) + elif not self.current_tool_name_sent: + normal_parts.append(content.strip()) + except (json.JSONDecodeError, ValueError, TypeError) as exc: + logger.warning("Failed to parse streamed dots tool call: %s", exc) + normal_parts.append(content.strip()) + + self.current_tool_name_sent = False + + return StreamingParseResult(normal_text="".join(normal_parts), calls=calls) + + def _stream_complete_json_body( + self, tools: list[Tool], calls: list[ToolCallItem] + ) -> None: + """Emit a complete JSON body while its closing XML tag is pending.""" + content = self._buffer[len(self.bot_token) :].strip() + if not content or not _is_complete_json(content): + return + + try: + parsed = json.loads(content) + except (json.JSONDecodeError, ValueError): + return + if not isinstance(parsed, dict): + return + + validated = self.parse_base_json(parsed, tools) + if not validated: + return + + item = validated[0] + arguments = item.parameters or "" + if not self.current_tool_name_sent: + self.current_tool_id += 1 + calls.append( + ToolCallItem( + tool_index=self.current_tool_id, + name=item.name, + parameters="", + ) + ) + self.prev_tool_call_arr.append( + {"name": item.name, "arguments": parsed.get("arguments", {})} + ) + self.streamed_args_for_tool.append("") + self.current_tool_name_sent = True + + streamed = self.streamed_args_for_tool[self.current_tool_id] + argument_diff = arguments.removeprefix(streamed) + if argument_diff: + calls.append( + ToolCallItem( + tool_index=self.current_tool_id, + name=None, + parameters=argument_diff, + ) + ) + self.streamed_args_for_tool[self.current_tool_id] += argument_diff + + def flush_pending_normal_text(self) -> str: + """Flush a partial opening marker as plain text at end of stream.""" + if not self._buffer or self.bot_token in self._buffer: + return "" + + normal_text = self._buffer.replace(self.eot_token, "") + self._buffer = "" + return normal_text + + def supports_structural_tag(self) -> bool: + return False + + def structure_info(self) -> _GetInfoFunc: + # Kept for the detector interface. It is not used while structural tags + # are disabled for dots' mixed XML/JSON format. + return lambda name: StructureInfo( + begin=f'{self.bot_token}{{"name": "{name}", "arguments": ', + end=f"}}{self.eot_token}", + trigger=self.bot_token, + ) diff --git a/python/sglang/srt/function_call/function_call_parser.py b/python/sglang/srt/function_call/function_call_parser.py index 011114b46..1530227b8 100644 --- a/python/sglang/srt/function_call/function_call_parser.py +++ b/python/sglang/srt/function_call/function_call_parser.py @@ -19,6 +19,7 @@ from sglang.srt.function_call.deepseekv3_detector import DeepSeekV3Detector from sglang.srt.function_call.deepseekv4_detector import DeepSeekV4Detector from sglang.srt.function_call.deepseekv31_detector import DeepSeekV31Detector from sglang.srt.function_call.deepseekv32_detector import DeepSeekV32Detector +from sglang.srt.function_call.dots_detector import DotsToolDetector from sglang.srt.function_call.gemma4_detector import Gemma4Detector from sglang.srt.function_call.gigachat3_detector import GigaChat3Detector from sglang.srt.function_call.glm4_moe_detector import Glm4MoeDetector @@ -68,6 +69,7 @@ class FunctionCallParser: "deepseekv31": DeepSeekV31Detector, "deepseekv32": DeepSeekV32Detector, "deepseekv4": DeepSeekV4Detector, + "dots": DotsToolDetector, "glm": Glm4MoeDetector, "glm45": Glm4MoeDetector, "glm47": Glm47MoeDetector, diff --git a/python/sglang/srt/layers/attention/attention_registry.py b/python/sglang/srt/layers/attention/attention_registry.py index 1d176ce7b..b1bd3c2ad 100644 --- a/python/sglang/srt/layers/attention/attention_registry.py +++ b/python/sglang/srt/layers/attention/attention_registry.py @@ -299,13 +299,28 @@ def attn_backend_wrapper_for_draft_extend( the mamba hybrids whose MTP draft is all softmax attention. Inkling's draft has its own short convs, so it must expose ``conv_state_metadata`` too. """ + from sglang.srt.configs.dots3 import Dots3Config from sglang.srt.configs.inkling import InklingMMConfig, InklingModelConfig if isinstance(runner.model_config.hf_config, (InklingModelConfig, InklingMMConfig)): return attn_backend_wrapper(runner, full_attn_backend) + if isinstance(runner.model_config.hf_text_config, Dots3Config): + return attn_backend_wrapper(runner, full_attn_backend) return full_attn_backend +def attn_backend_wrapper_for_draft_decode(runner: "ModelRunner", backend): + """Apply the Dots model wrapper to per-step draft backends.""" + from sglang.srt.configs.dots3 import Dots3Config + + if not hasattr(runner, "model_config"): + return backend + hf_text_config = runner.model_config.hf_text_config + if isinstance(hf_text_config, Dots3Config): + return hf_text_config.wrap_draft_decode_attention_backend(backend) + return backend + + def attn_backend_wrapper(runner: "ModelRunner", full_attn_backend: "AttentionBackend"): """ Wrapper for special models like hybrid GDN, so we don't @@ -315,8 +330,14 @@ def attn_backend_wrapper(runner: "ModelRunner", full_attn_backend: "AttentionBac hybrid_gdn_config(runner.model_config) is not None and runner.use_mla_backend ), "hybrid_gdn can only be used with non-MLA models." + from sglang.srt.configs.dots3 import Dots3Config from sglang.srt.configs.model_config import is_minimax_sparse + if isinstance(runner.model_config.hf_text_config, Dots3Config): + return runner.model_config.hf_text_config.wrap_attention_backend( + runner, full_attn_backend + ) + if is_minimax_sparse(runner.model_config.hf_config): from sglang.srt.layers.attention.minimax_sparse_backend import ( MiniMaxHybridAttnBackend, diff --git a/python/sglang/srt/layers/attention/dots_hybrid_backend.py b/python/sglang/srt/layers/attention/dots_hybrid_backend.py new file mode 100644 index 000000000..887f70b4e --- /dev/null +++ b/python/sglang/srt/layers/attention/dots_hybrid_backend.py @@ -0,0 +1,571 @@ +"""Layer-wise DSA/SWA attention dispatch for dots.note.omni.""" + +from __future__ import annotations + +from contextlib import contextmanager +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import torch + +from sglang.kernels.ops.attention.flash_attention import flash_attn_varlen_func +from sglang.srt.layers.attention.base_attn_backend import ( + AttentionBackend, + SharedReadEnds, +) +from sglang.srt.layers.attention.hybrid_attn_backend import HybridAttnBackend + +if TYPE_CHECKING: + from sglang.srt.layers.radix_attention import RadixAttention + from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode + from sglang.srt.speculative.spec_info import SpecInput + + +def _normalize_page_table_rows( + page_table: torch.Tensor, batch_size: int +) -> torch.Tensor: + """Match Dots' pre-planned SWA table to the live DP-padded batch.""" + if page_table.shape[0] >= batch_size: + return page_table[:batch_size] + return torch.cat( + [ + page_table, + page_table.new_zeros( + (batch_size - page_table.shape[0], page_table.shape[1]) + ), + ], + dim=0, + ) + + +def _normalize_cache_seqlens_rows( + cache_seqlens: torch.Tensor, + seq_lens: torch.Tensor, + batch_size: int, +) -> torch.Tensor: + """Preserve planned rows and fill only newly DP-padded dummy rows.""" + planned_bs = cache_seqlens.shape[0] + if planned_bs >= batch_size: + return cache_seqlens[:batch_size] + + dummy_seqlens = seq_lens[planned_bs:batch_size].to( + device=cache_seqlens.device, + dtype=cache_seqlens.dtype, + non_blocking=True, + ) + return torch.cat([cache_seqlens, dummy_seqlens], dim=0) + + +def _metadata_mismatches_dp_padded_batch(metadata, forward_batch) -> bool: + """True when pre-planned attention metadata no longer matches the live batch. + + EAGLE plans draft metadata before ModelRunner runs DP/MLP padding. Dummy + request rows then change ``batch_size`` / ``out_cache_loc``, leaving + page tables and SWA write targets short. Rebuilding is required; slicing + or zero-padding the stale tensors is not enough for DSA + SWA. + """ + if metadata is None: + return False + bs = forward_batch.batch_size + from sglang.srt.layers.attention.flashattention_backend import ( + FlashAttentionMetadata, + ) + + if isinstance(metadata, FlashAttentionMetadata): + if metadata.page_table is not None and metadata.page_table.shape[0] != bs: + return True + if ( + metadata.swa_page_table is not None + and metadata.swa_page_table.shape[0] != bs + ): + return True + if ( + metadata.cache_seqlens_int32 is not None + and metadata.cache_seqlens_int32.shape[0] != bs + ): + return True + swa_loc = metadata.swa_out_cache_loc + out_loc = forward_batch.out_cache_loc + return ( + swa_loc is not None + and out_loc is not None + and swa_loc.shape[0] != out_loc.shape[0] + ) + + from sglang.srt.layers.attention.dsa_backend import DSAMetadata + + if isinstance(metadata, DSAMetadata): + return metadata.cache_seqlens_int32.shape[0] != bs + return False + + +def _dp_padding_changed_batch_size(forward_batch) -> bool: + original_bs = forward_batch._original_batch_size + return original_bs is not None and original_bs != forward_batch.batch_size + + +def _maybe_rebuild_dots_metadata(backend, forward_batch) -> None: + """Eager-only rebuild when DP padding invalidated a Dots pre-plan.""" + from sglang.srt.model_executor.runner_utils.capture_mode import ( + get_is_capture_mode, + ) + + if get_is_capture_mode(): + return + if backend._dp_rebuilt_batch_id == id(forward_batch): + return + stale = _metadata_mismatches_dp_padded_batch( + backend.forward_metadata, forward_batch + ) or _dp_padding_changed_batch_size(forward_batch) + if not stale: + return + backend.init_forward_metadata(forward_batch) + backend._dp_rebuilt_batch_id = id(forward_batch) + + +@dataclass +class DotsSWAMLAPrefillMetadata: + kv_indices: torch.Tensor + cu_seqlens_q: torch.Tensor + cu_seqlens_k: torch.Tensor + max_seq_len_q: int + max_seq_len_k: int + + +class DotsSWAMLAAttnBackend(AttentionBackend): + """Add Dots latent-cache SWA support around a FlashAttention backend.""" + + def __init__(self, backend: AttentionBackend): + self.backend = backend + self._active_backend = backend + self.token_to_kv_pool = backend.token_to_kv_pool + self.req_to_token_pool = backend.req_to_token_pool + self.needs_cpu_seq_lens = True + self._prefill_metadata: DotsSWAMLAPrefillMetadata | None = None + self._dp_rebuilt_batch_id: int | None = None + + @property + def forward_metadata(self): + return self._active_backend.forward_metadata + + @forward_metadata.setter + def forward_metadata(self, value): + self._active_backend.forward_metadata = value + + @property + def verify_mask(self): + return self.backend.verify_mask + + def shared_read_ends(self, fm: ForwardMode) -> SharedReadEnds: + return self.backend.shared_read_ends(fm) + + def draft_extend_metadata_captured_in_graph(self) -> bool: + return self.backend.draft_extend_metadata_captured_in_graph() + + def selected_backend(self, forward_batch: ForwardBatch) -> AttentionBackend: + return ( + self.backend._select_backend(forward_batch.forward_mode) + if isinstance(self.backend, HybridAttnBackend) + else self.backend + ) + + def uses_flash_attention(self, forward_batch: ForwardBatch) -> bool: + from sglang.srt.layers.attention.flashattention_backend import ( + FlashAttentionBackend, + ) + + return isinstance(self.selected_backend(forward_batch), FlashAttentionBackend) + + def maybe_rebuild_metadata_after_dp_padding( + self, forward_batch: ForwardBatch + ) -> None: + """Rebuild FA + SWA-prefill metadata after eager DP dummy-row padding.""" + self._active_backend = self.selected_backend(forward_batch) + _maybe_rebuild_dots_metadata(self, forward_batch) + + def select_draft_step_out_cache_loc(self, forward_batch: ForwardBatch): + """Return this draft step's write locations from a combined SWA buffer.""" + from sglang.srt.layers.attention.flashattention_backend import ( + FlashAttentionBackend, + ) + + out_cache_loc = forward_batch.out_cache_loc + backend = self._active_backend + if not isinstance(backend, FlashAttentionBackend): + return out_cache_loc + if ( + out_cache_loc is not None + and forward_batch.forward_mode.is_decode_or_idle() + and forward_batch.spec_info is not None + and backend.speculative_num_steps > 0 + and out_cache_loc.numel() + == forward_batch.batch_size * backend.topk * backend.speculative_num_steps + ): + return out_cache_loc.view( + forward_batch.batch_size, + backend.topk, + backend.speculative_num_steps, + )[:, :, backend.speculative_step_id].reshape(-1) + return out_cache_loc + + @contextmanager + def _use_draft_step_out_cache_loc(self, forward_batch: ForwardBatch): + """Expose only this backend's draft-step write locations to FA.""" + original = forward_batch.out_cache_loc + forward_batch.out_cache_loc = self.select_draft_step_out_cache_loc( + forward_batch + ) + try: + yield + finally: + forward_batch.out_cache_loc = original + + def init_forward_metadata(self, forward_batch: ForwardBatch): + self._active_backend = self.selected_backend(forward_batch) + with self._use_draft_step_out_cache_loc(forward_batch): + self.backend.init_forward_metadata(forward_batch) + self._init_prefill_metadata(forward_batch) + + def init_forward_metadata_out_graph( + self, forward_batch: ForwardBatch, in_capture: bool = False + ): + self._active_backend = self.selected_backend(forward_batch) + with self._use_draft_step_out_cache_loc(forward_batch): + self.backend.init_forward_metadata_out_graph( + forward_batch, in_capture=in_capture + ) + self._init_prefill_metadata(forward_batch) + + def init_forward_metadata_in_graph(self, forward_batch: ForwardBatch): + self.backend.init_forward_metadata_in_graph(forward_batch) + + def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int): + self.backend.init_cuda_graph_state(max_bs, max_num_tokens) + + def get_cuda_graph_seq_len_fill_value(self): + return self.backend.get_cuda_graph_seq_len_fill_value() + + def on_after_cuda_graph_warmup(self): + self.backend.on_after_cuda_graph_warmup() + + def update_verify_buffers_to_fill_after_draft( + self, spec_info: SpecInput, cuda_graph_bs: int | None + ): + return self.backend.update_verify_buffers_to_fill_after_draft( + spec_info, cuda_graph_bs + ) + + def forward(self, q, k, v, layer, forward_batch, save_kv_cache=True, **kwargs): + self.maybe_rebuild_metadata_after_dp_padding(forward_batch) + return self.backend.forward( + q, k, v, layer, forward_batch, save_kv_cache, **kwargs + ) + + def forward_extend( + self, q, k, v, layer, forward_batch, save_kv_cache=True, **kwargs + ): + return self.backend.forward_extend( + q, k, v, layer, forward_batch, save_kv_cache, **kwargs + ) + + def forward_decode( + self, q, k, v, layer, forward_batch, save_kv_cache=True, **kwargs + ): + return self.backend.forward_decode( + q, k, v, layer, forward_batch, save_kv_cache, **kwargs + ) + + def init_mha_chunk_metadata(self, forward_batch: ForwardBatch): + self.backend.init_mha_chunk_metadata(forward_batch) + + def _init_prefill_metadata(self, forward_batch: ForwardBatch): + if not forward_batch.forward_mode.is_extend_without_speculative(): + self._prefill_metadata = None + return + + metadata = self._active_backend.forward_metadata + assert forward_batch.seq_lens_cpu is not None + batch_kv_indices = self._active_backend.req_to_token[ + forward_batch.req_pool_indices, : + ] + sliced_indices = [] + kv_lens = [] + for i in range(forward_batch.batch_size): + q_len = int(forward_batch.extend_seq_lens_cpu[i]) + kv_len = int(forward_batch.seq_lens_cpu[i]) + tail_len = min(q_len + self._active_backend.sliding_window_size, kv_len) + sliced_indices.append(batch_kv_indices[i, kv_len - tail_len : kv_len]) + kv_lens.append(tail_len) + + full_kv_indices = torch.cat(sliced_indices) + kv_indices = self.token_to_kv_pool.translate_loc_from_full_to_swa( + full_kv_indices + ).to(torch.int32) + lens_cpu = torch.tensor([0, *kv_lens], dtype=torch.int32, pin_memory=True) + self._prefill_metadata = DotsSWAMLAPrefillMetadata( + kv_indices=kv_indices, + cu_seqlens_q=metadata.cu_seqlens_q, + cu_seqlens_k=torch.cumsum( + lens_cpu.to(device=forward_batch.seq_lens.device, non_blocking=True), + dim=0, + dtype=torch.int32, + ), + max_seq_len_q=metadata.max_seq_len_q, + max_seq_len_k=max(kv_lens), + ) + + def get_swa_mla_prefill_latent_cache( + self, forward_batch: ForwardBatch, layer_id: int + ): + assert self._prefill_metadata is not None + return self.token_to_kv_pool.get_key_buffer(layer_id)[ + self._prefill_metadata.kv_indices + ] + + def forward_swa_mla_expanded(self, q, k, v, layer, forward_batch=None): + """Run dense SWA after Dots expands its compact MLA cache.""" + metadata = self._prefill_metadata + assert metadata is not None + q = q.view(-1, layer.tp_q_head_num, layer.head_dim) + k = k.view(-1, layer.tp_k_head_num, layer.head_dim).to(q.dtype) + v = v.view(-1, layer.tp_k_head_num, layer.v_head_dim).to(q.dtype) + + # FA3 requires equal QK/V widths when QK exceeds 192. + pad_v_to_qk = layer.head_dim > 192 and layer.v_head_dim != layer.head_dim + if pad_v_to_qk: + v = torch.nn.functional.pad(v, (0, layer.head_dim - layer.v_head_dim)) + + output = flash_attn_varlen_func( + q=q, + k=k, + v=v, + cu_seqlens_q=metadata.cu_seqlens_q, + cu_seqlens_k=metadata.cu_seqlens_k, + max_seqlen_q=metadata.max_seq_len_q, + max_seqlen_k=metadata.max_seq_len_k, + softmax_scale=layer.scaling, + causal=True, + window_size=(layer.sliding_window_size, 0), + ver=self._active_backend.fa_impl_ver, + ) + if pad_v_to_qk: + output = output[..., : layer.v_head_dim] + return output.reshape(-1, layer.tp_q_head_num * layer.v_head_dim) + + def forward_swa_mla_absorbed(self, q, layer, forward_batch): + """Run decode directly against the page64 latent SWA cache.""" + from sglang.srt.layers.attention.swa_mla_fallback.forward import ( + forward_dense_kvlora_swa_torch_fallback, + ) + + backend = self.selected_backend(forward_batch) + if backend.page_size != 64: + raise RuntimeError( + "Dots SWA latent decode requires page_size=64, " + f"got {backend.page_size}." + ) + + self.maybe_rebuild_metadata_after_dp_padding(forward_batch) + metadata = backend.forward_metadata + block_table = metadata.swa_page_table + if block_table is None: + raise RuntimeError("Dots SWA latent decode requires an SWA page table.") + bs = forward_batch.batch_size + block_table = _normalize_page_table_rows(block_table, bs) + cache_seqlens = _normalize_cache_seqlens_rows( + metadata.cache_seqlens_int32, + forward_batch.seq_lens, + bs, + ) + reshape_q = q.view(bs, -1, layer.tp_q_head_num, layer.head_dim) + k_cache = self.token_to_kv_pool.get_key_buffer(layer.layer_id) + output = forward_dense_kvlora_swa_torch_fallback( + reshape_q=reshape_q, + k_cache=k_cache, + block_table=block_table, + cache_seqlens=cache_seqlens, + layer=layer, + kv_cache_dim=layer.head_dim, + head_dim_v=layer.v_head_dim, + window_size=layer.sliding_window_size + 1, + ) + return output.view(-1, layer.tp_q_head_num * layer.v_head_dim) + + +class DotsHybridAttnBackend(AttentionBackend): + def __init__( + self, + dsa_backend: AttentionBackend, + swa_backend: AttentionBackend, + ): + self.dsa_backend = dsa_backend + # Keep DSA on its radix-aware MLA path. + self.dsa_backend.supports_mha_one_shot = False + self.swa_backend = swa_backend + self.token_to_kv_pool = swa_backend.token_to_kv_pool + self.req_to_token_pool = swa_backend.req_to_token_pool + # SWA latent expansion uses host sequence-length mirrors. + self.needs_cpu_seq_lens = True + self._dp_rebuilt_batch_id: int | None = None + + @staticmethod + def _is_swa_layer(layer: RadixAttention) -> bool: + return layer.sliding_window_size is not None and layer.sliding_window_size > -1 + + def backend_for_layer(self, layer: RadixAttention) -> AttentionBackend: + return self.swa_backend if self._is_swa_layer(layer) else self.dsa_backend + + def selected_swa_backend(self, forward_batch: ForwardBatch) -> AttentionBackend: + return ( + self.swa_backend._select_backend(forward_batch.forward_mode) + if isinstance(self.swa_backend, HybridAttnBackend) + else self.swa_backend + ) + + def maybe_rebuild_metadata_after_dp_padding( + self, forward_batch: ForwardBatch + ) -> None: + """Rebuild both DSA and SWA plans after eager DP dummy-row padding.""" + from sglang.srt.model_executor.runner_utils.capture_mode import ( + get_is_capture_mode, + ) + + if get_is_capture_mode(): + return + if self._dp_rebuilt_batch_id == id(forward_batch): + return + dsa_stale = _metadata_mismatches_dp_padded_batch( + self.dsa_backend.forward_metadata, forward_batch + ) + swa_backend = self.selected_swa_backend(forward_batch) + swa_stale = _metadata_mismatches_dp_padded_batch( + swa_backend.forward_metadata, forward_batch + ) + if dsa_stale or swa_stale or _dp_padding_changed_batch_size(forward_batch): + self.init_forward_metadata(forward_batch) + self._dp_rebuilt_batch_id = id(forward_batch) + + def init_forward_metadata(self, forward_batch: ForwardBatch): + self.dsa_backend.init_forward_metadata(forward_batch) + self.swa_backend.init_forward_metadata(forward_batch) + + def init_forward_metadata_out_graph( + self, forward_batch: ForwardBatch, in_capture: bool = False + ): + self.dsa_backend.init_forward_metadata_out_graph( + forward_batch, in_capture=in_capture + ) + self.swa_backend.init_forward_metadata_out_graph( + forward_batch, in_capture=in_capture + ) + + def init_forward_metadata_in_graph(self, forward_batch: ForwardBatch): + self.dsa_backend.init_forward_metadata_in_graph(forward_batch) + self.swa_backend.init_forward_metadata_in_graph(forward_batch) + + def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int): + self.dsa_backend.init_cuda_graph_state(max_bs, max_num_tokens) + self.swa_backend.init_cuda_graph_state(max_bs, max_num_tokens) + + def get_cuda_graph_seq_len_fill_value(self): + return self.swa_backend.get_cuda_graph_seq_len_fill_value() + + def on_after_cuda_graph_warmup(self): + self.dsa_backend.on_after_cuda_graph_warmup() + self.swa_backend.on_after_cuda_graph_warmup() + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: RadixAttention, + forward_batch: ForwardBatch, + save_kv_cache: bool = True, + **kwargs, + ): + self.maybe_rebuild_metadata_after_dp_padding(forward_batch) + return self.backend_for_layer(layer).forward( + q, k, v, layer, forward_batch, save_kv_cache, **kwargs + ) + + def forward_extend( + self, q, k, v, layer, forward_batch, save_kv_cache=True, **kwargs + ): + return self.backend_for_layer(layer).forward_extend( + q, k, v, layer, forward_batch, save_kv_cache, **kwargs + ) + + def forward_decode( + self, q, k, v, layer, forward_batch, save_kv_cache=True, **kwargs + ): + return self.backend_for_layer(layer).forward_decode( + q, k, v, layer, forward_batch, save_kv_cache, **kwargs + ) + + def get_indexer_metadata(self, layer_id: int, forward_batch: ForwardBatch): + return self.dsa_backend.get_indexer_metadata(layer_id, forward_batch) + + def get_swa_mla_prefill_latent_cache( + self, forward_batch: ForwardBatch, layer_id: int + ): + backend = self.selected_swa_backend(forward_batch) + return backend.get_swa_mla_prefill_latent_cache(forward_batch, layer_id) + + def forward_swa_mla_expanded(self, q, k, v, layer, forward_batch): + backend = self.selected_swa_backend(forward_batch) + return backend.forward_swa_mla_expanded(q, k, v, layer, forward_batch) + + def forward_swa_mla_absorbed(self, q, layer, forward_batch): + backend = self.selected_swa_backend(forward_batch) + return backend.forward_swa_mla_absorbed(q, layer, forward_batch) + + def init_mha_chunk_metadata(self, forward_batch: ForwardBatch): + backend = self.selected_swa_backend(forward_batch) + backend.init_mha_chunk_metadata(forward_batch) + + +def _wrap_dots_swa_backend(backend: AttentionBackend) -> AttentionBackend: + """Add latent-cache SWA behavior when a backend uses FlashAttention.""" + from sglang.srt.layers.attention.flashattention_backend import ( + FlashAttentionBackend, + ) + + if isinstance(backend, FlashAttentionBackend) or ( + isinstance(backend, HybridAttnBackend) + and ( + isinstance(backend.prefill_backend, FlashAttentionBackend) + or isinstance(backend.decode_backend, FlashAttentionBackend) + ) + ): + return DotsSWAMLAAttnBackend(backend) + return backend + + +def wrap_dots_draft_decode_backend(backend: AttentionBackend) -> AttentionBackend: + """Wrap each per-step backend used by the Dots NextN draft container.""" + backend.attn_backends = [ + _wrap_dots_swa_backend(child) for child in backend.attn_backends + ] + return backend + + +def wrap_dots_attention_backend(runner, full_attn_backend: AttentionBackend): + """Construct the Dots target or draft attention backend.""" + if runner.model_config.is_draft_model: + return _wrap_dots_swa_backend(full_attn_backend) + + if runner.model_config.hf_text_config.index_topk is None: + return DotsSWAMLAAttnBackend(full_attn_backend) + + from sglang.srt.layers.attention.attention_registry import create_dsa_backend + + swa_backend = ( + full_attn_backend.prefill_backend + if isinstance(full_attn_backend, HybridAttnBackend) + else full_attn_backend + ) + return DotsHybridAttnBackend( + dsa_backend=create_dsa_backend(runner), + swa_backend=DotsSWAMLAAttnBackend(swa_backend), + ) diff --git a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py index eaa5cc604..42a0c2208 100644 --- a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py +++ b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py @@ -575,7 +575,13 @@ class Indexer(DSANPUIndexerMixin, BaseFusedOp): key, [self.rope_head_dim, self.head_dim - self.rope_head_dim], dim=-1 ) - _, k_rope = self.rotary_emb(positions, k_rope, k_rope) + # Rotary may update both inputs in place, so the K-only path must not + # alias its dummy query with the key. + if _is_cuda or _is_hip or _is_xpu: + dummy_q_rope = torch.empty_like(k_rope) + else: + dummy_q_rope = k_rope + _, k_rope = self.rotary_emb(positions, dummy_q_rope, k_rope) self._update_rope_guarded(key[..., : self.rope_head_dim], k_rope) key = rotate_activation(key) diff --git a/python/sglang/srt/layers/attention/dsa_backend.py b/python/sglang/srt/layers/attention/dsa_backend.py index 9078dbac6..41789a03c 100644 --- a/python/sglang/srt/layers/attention/dsa_backend.py +++ b/python/sglang/srt/layers/attention/dsa_backend.py @@ -333,6 +333,7 @@ class DeepseekSparseAttnBackend( self.req_to_token = model_runner.req_to_token_pool.req_to_token self.use_mha: bool = False + self.supports_mha_one_shot: bool = True self.dsa_prefill_impl: _DSA_IMPL_T = ( model_runner.server_args.dsa_prefill_backend ) @@ -3325,7 +3326,8 @@ class DeepseekSparseAttnBackend( # Requirements: H200/B200/MI355X, short sequences, supported dtype, fits in chunk self.use_mha = ( - ( + self.supports_mha_one_shot + and ( device_sm == 90 or (device_sm >= 100 and device_sm < 110) or _IS_GFX95 diff --git a/python/sglang/srt/layers/attention/swa_mla_fallback/__init__.py b/python/sglang/srt/layers/attention/swa_mla_fallback/__init__.py new file mode 100644 index 000000000..4d88759b2 --- /dev/null +++ b/python/sglang/srt/layers/attention/swa_mla_fallback/__init__.py @@ -0,0 +1 @@ +"""Fallback operations for sliding-window MLA attention paths.""" diff --git a/python/sglang/srt/layers/attention/swa_mla_fallback/forward.py b/python/sglang/srt/layers/attention/swa_mla_fallback/forward.py new file mode 100644 index 000000000..26570a339 --- /dev/null +++ b/python/sglang/srt/layers/attention/swa_mla_fallback/forward.py @@ -0,0 +1,77 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from sglang.srt.layers.attention.swa_mla_fallback.ops import ( + apply_swa_score_mask, + gather_page64_kv_latent, +) + +if TYPE_CHECKING: + from sglang.srt.layers.radix_attention import RadixAttention + + +def forward_dense_kvlora_swa_torch_fallback( + reshape_q: torch.Tensor, + k_cache: torch.Tensor, + block_table: torch.Tensor, + cache_seqlens: torch.Tensor, + layer: RadixAttention, + kv_cache_dim: int, + head_dim_v: int, + window_size: int, +) -> torch.Tensor: + """Page-64 SWA fallback for dense KV-LoRA decode.""" + if layer.tp_k_head_num != 1: + raise RuntimeError( + "SWA MLA torch fallback currently supports MLA with one " + f"KV head, got tp_k_head_num={layer.tp_k_head_num}." + ) + + bs, s_q, num_heads, qk_dim = reshape_q.shape + if qk_dim != kv_cache_dim: + raise RuntimeError( + f"SWA MLA torch fallback got q dim {qk_dim}, " + f"expected kv_cache_dim {kv_cache_dim}." + ) + if s_q not in (1, 4): + raise RuntimeError( + "SWA MLA torch fallback mask is specialized for s_q=1 " + f"or s_q=4, got s_q={s_q}." + ) + + # Include the full union of causal windows and align it for BMM. + kv_latent, kv_valid = gather_page64_kv_latent( + k_cache, + block_table, + cache_seqlens, + window_size, + s_q, + kv_cache_dim, + ) + gather_len = kv_latent.shape[1] + + # Keep the output in [bs, s_q, num_heads, head_dim_v] order. + q_for_scores = reshape_q.reshape(bs, s_q * num_heads, qk_dim) + scores = torch.bmm(q_for_scores, kv_latent.transpose(1, 2)).view( + bs, s_q, num_heads, gather_len + ) + scores = scores.float() + scores.mul_(layer.scaling) + + apply_swa_score_mask( + scores.transpose(1, 2), + cache_seqlens, + kv_valid, + num_heads, + window_size, + s_q, + ) + + probs = torch.softmax(scores, dim=-1).to(reshape_q.dtype) + return torch.bmm( + probs.reshape(bs, s_q * num_heads, gather_len), + kv_latent[..., :head_dim_v], + ).view(bs, s_q, num_heads, head_dim_v) diff --git a/python/sglang/srt/layers/attention/swa_mla_fallback/ops.py b/python/sglang/srt/layers/attention/swa_mla_fallback/ops.py new file mode 100644 index 000000000..0d12155af --- /dev/null +++ b/python/sglang/srt/layers/attention/swa_mla_fallback/ops.py @@ -0,0 +1,205 @@ +import torch +import triton +import triton.language as tl + +PAGE_SIZE = 64 + + +@triton.jit +def _gather_page64_kv_latent_kernel( + k_cache_ptr, + block_table_ptr, + cache_seqlens_ptr, + kv_out_ptr, + valid_out_ptr, + k_cache_stride_t, + k_cache_stride_d, + block_table_stride_b, + block_table_num_pages: tl.constexpr, + k_cache_num_tokens: tl.constexpr, + kv_out_stride_b, + kv_out_stride_t, + kv_out_stride_d, + valid_out_stride_b, + valid_out_stride_t, + GATHER_LEN: tl.constexpr, + KV_DIM: tl.constexpr, + BLOCK_T: tl.constexpr, + BLOCK_D: tl.constexpr, + PAGE_SIZE_: tl.constexpr, +): + bid = tl.program_id(0) + tid = tl.program_id(1) + did = tl.program_id(2) + + cache_seqlen = tl.load(cache_seqlens_ptr + bid).to(tl.int32) + gather_start = tl.maximum(cache_seqlen - GATHER_LEN, 0) + offs_t = tid * BLOCK_T + tl.arange(0, BLOCK_T) + logical_token = gather_start + offs_t + logical_valid = (offs_t < GATHER_LEN) & (logical_token < cache_seqlen) + + logical_page = logical_token // PAGE_SIZE_ + intra_page = logical_token - logical_page * PAGE_SIZE_ + page_table_valid = logical_valid & (logical_page < block_table_num_pages) + physical_page = tl.load( + block_table_ptr + bid * block_table_stride_b + logical_page, + mask=page_table_valid, + other=-1, + ).to(tl.int32) + physical_token = physical_page * PAGE_SIZE_ + intra_page + physical_valid = ( + page_table_valid & (physical_page >= 0) & (physical_token < k_cache_num_tokens) + ) + + offs_d = did * BLOCK_D + tl.arange(0, BLOCK_D) + values = tl.load( + k_cache_ptr + + physical_token[:, None] * k_cache_stride_t + + offs_d[None, :] * k_cache_stride_d, + mask=physical_valid[:, None] & (offs_d[None, :] < KV_DIM), + other=0.0, + ) + tl.store( + kv_out_ptr + + bid * kv_out_stride_b + + offs_t[:, None] * kv_out_stride_t + + offs_d[None, :] * kv_out_stride_d, + values, + mask=(offs_t[:, None] < GATHER_LEN) & (offs_d[None, :] < KV_DIM), + ) + tl.store( + valid_out_ptr + bid * valid_out_stride_b + offs_t * valid_out_stride_t, + physical_valid, + mask=(did == 0) & (offs_t < GATHER_LEN), + ) + + +def gather_page64_kv_latent( + k_cache: torch.Tensor, + block_table: torch.Tensor, + cache_seqlens: torch.Tensor, + window_size: int, + s_q: int, + kv_cache_dim: int, +): + bs = cache_seqlens.shape[0] + assert block_table.shape[0] == bs + + gather_len = ((window_size + s_q - 1 + 7) // 8) * 8 + kv_latent = torch.empty( + (bs, gather_len, kv_cache_dim), + dtype=k_cache.dtype, + device=k_cache.device, + ) + kv_valid = torch.empty((bs, gather_len), dtype=torch.bool, device=k_cache.device) + + block_t = 8 + block_d = 128 + _gather_page64_kv_latent_kernel[ + (bs, triton.cdiv(gather_len, block_t), triton.cdiv(kv_cache_dim, block_d)) + ]( + k_cache, + block_table, + cache_seqlens, + kv_latent, + kv_valid, + k_cache.stride(0), + k_cache.stride(2), + block_table.stride(0), + block_table.shape[1], + k_cache.shape[0], + kv_latent.stride(0), + kv_latent.stride(1), + kv_latent.stride(2), + kv_valid.stride(0), + kv_valid.stride(1), + GATHER_LEN=gather_len, + KV_DIM=kv_cache_dim, + BLOCK_T=block_t, + BLOCK_D=block_d, + PAGE_SIZE_=PAGE_SIZE, + num_warps=4, + ) + return kv_latent, kv_valid + + +@triton.jit +def _apply_swa_score_mask_kernel( + scores_ptr, + cache_seqlens_ptr, + valid_ptr, + scores_stride_b, + scores_stride_h, + scores_stride_q, + scores_stride_t, + valid_stride_b, + valid_stride_t, + GATHER_LEN: tl.constexpr, + WINDOW_SIZE: tl.constexpr, + S_Q: tl.constexpr, + BLOCK_T: tl.constexpr, +): + bid = tl.program_id(0) + hid = tl.program_id(1) + qid = tl.program_id(2) + + offs_t = tl.arange(0, BLOCK_T) + cache_seqlen = tl.load(cache_seqlens_ptr + bid).to(tl.int32) + gather_start = tl.maximum(cache_seqlen - GATHER_LEN, 0) + kv_pos = gather_start + offs_t + q_pos = cache_seqlen - S_Q + qid + page_valid = tl.load( + valid_ptr + bid * valid_stride_b + offs_t * valid_stride_t, + mask=offs_t < GATHER_LEN, + other=0, + ).to(tl.int1) + valid = ( + (offs_t < GATHER_LEN) + & page_valid + & (kv_pos <= q_pos) + & (kv_pos >= q_pos - WINDOW_SIZE + 1) + & (q_pos >= 0) + ) + tl.store( + scores_ptr + + bid * scores_stride_b + + hid * scores_stride_h + + qid * scores_stride_q + + offs_t * scores_stride_t, + tl.full((BLOCK_T,), -3.4028234663852886e38, tl.float32), + mask=(offs_t < GATHER_LEN) & ~valid, + ) + + +def apply_swa_score_mask( + scores: torch.Tensor, + cache_seqlens: torch.Tensor, + kv_valid: torch.Tensor, + num_heads: int, + window_size: int, + s_q: int, +): + bs = cache_seqlens.shape[0] + assert scores.shape[0] == bs + assert scores.shape[1] == num_heads + assert scores.shape[2] == s_q + assert kv_valid.shape == (bs, scores.shape[3]) + gather_len = scores.shape[3] + mask_block_t = triton.next_power_of_2(gather_len) + _apply_swa_score_mask_kernel[(bs, num_heads, s_q)]( + scores, + cache_seqlens, + kv_valid, + scores.stride(0), + scores.stride(1), + scores.stride(2), + scores.stride(3), + kv_valid.stride(0), + kv_valid.stride(1), + GATHER_LEN=gather_len, + WINDOW_SIZE=window_size, + S_Q=s_q, + BLOCK_T=mask_block_t, + num_warps=4, + ) + return scores diff --git a/python/sglang/srt/layers/quantization/base_config.py b/python/sglang/srt/layers/quantization/base_config.py index f702bd8f8..186e7f9b0 100644 --- a/python/sglang/srt/layers/quantization/base_config.py +++ b/python/sglang/srt/layers/quantization/base_config.py @@ -128,6 +128,8 @@ class FusedMoEMethodBase(QuantizeMethodBase): class QuantizationConfig(ABC): """Base class for quantization configs.""" + weight_block_size: Optional[List[int]] = None + def __init__(self): super().__init__() # mapping is updated by models as they initialize diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index ca060faf9..f9f2f7df2 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -206,6 +206,8 @@ class GenerateReqInput: ] = None # Whether to extract and process audio from video inputs. use_audio_in_video: bool = False + # Optional request-scoped video processor configuration. + video_config: Optional[Dict[str, Any]] = None # The sampling_params. See descriptions below. sampling_params: Optional[Union[List[Dict[str, Any]], Dict[str, Any]]] = None # Whether to return logprobs. diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index e92d34d6d..4d366cf10 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -738,6 +738,8 @@ class KVCacheConfigurator: get_exec().kernel.attention_backend == "ascend" and not self.mambaish_config ): unsupported_pool_family = "NPU/Ascend KV pool" + elif self.use_mla_backend and self.is_hybrid_swa: + unsupported_pool_family = "hybrid DSA/MLA-SWA KV pool" elif self.use_mla_backend and is_dsa_model: unsupported_pool_family = "DSA/MLA KV pool" elif self.use_mla_backend and not self.mambaish_config: @@ -1013,6 +1015,12 @@ class KVCacheConfigurator: token_to_kv_pool = self._build_ascend_mha_kv_pool( max_total_num_tokens=sizes.max_total_num_tokens, ) + elif self.use_mla_backend and self.is_hybrid_swa: + token_to_kv_pool = self._build_hybrid_mla_swa_kv_pool( + full_max_total_num_tokens=sizes.full_max_total_num_tokens, + swa_max_total_num_tokens=sizes.swa_max_total_num_tokens, + is_dsa_model=is_dsa_model, + ) elif self.use_mla_backend and is_dsa_model: token_to_kv_pool = self._build_dsa_kv_pool( max_total_num_tokens=sizes.max_total_num_tokens, @@ -1358,6 +1366,60 @@ class KVCacheConfigurator: ) return token_to_kv_pool + def _build_hybrid_mla_swa_kv_pool( + self, + *, + full_max_total_num_tokens: int, + swa_max_total_num_tokens: int, + is_dsa_model: bool, + ) -> KVCache: + """Build a hybrid MLA pool with independent full/SWA cache geometries. + + Full-attention layers may use either MLA or DSA storage, while sliding + layers use MLA storage. The returned ``SWAKVPool`` exposes the common + MLA and optional DSA-index interfaces independent of model type. + """ + full_pool_class = DSATokenToKVPool if is_dsa_model else MLATokenToKVPool + common = { + "page_size": self.server_args.page_size, + "device": self.device, + "enable_memory_saver": False, + } + full_pool_kwargs = { + **common, + "kv_lora_rank": self.model_config.kv_lora_rank, + "qk_rope_head_dim": self.model_config.qk_rope_head_dim, + } + if is_dsa_model: + full_pool_kwargs.update( + index_head_dim=get_dsa_index_head_dim(self.model_config.hf_config), + kv_cache_dim=calculate_mla_kv_cache_dim( + model_config=self.model_config, + kv_cache_dtype=self.kv_cache_dtype, + server_args=self.server_args, + ), + ) + + return SWAKVPool( + size=full_max_total_num_tokens, + size_swa=swa_max_total_num_tokens, + page_size=self.server_args.page_size, + dtype=self.kv_cache_dtype, + head_num=0, + head_dim=0, + swa_attention_layer_ids=self.model_config.swa_attention_layer_ids, + full_attention_layer_ids=self.model_config.full_attention_layer_ids, + device=self.device, + full_kv_pool_class=full_pool_class, + swa_kv_pool_class=MLATokenToKVPool, + full_kv_pool_kwargs=full_pool_kwargs, + swa_kv_pool_kwargs={ + **common, + "kv_lora_rank": self.model_config.swa_kv_lora_rank, + "qk_rope_head_dim": self.model_config.swa_qk_rope_head_dim, + }, + ) + def _build_mla_fp4_kv_pool(self, *, max_total_num_tokens: int) -> KVCache: token_to_kv_pool = MLATokenToKVPoolFP4( max_total_num_tokens, diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index da331decf..43f4f26d7 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -4125,10 +4125,13 @@ class MLATokenToKVPool(KVCache): loc_info, cache_k: torch.Tensor, cache_v: torch.Tensor, + layer_id_override: Optional[int] = None, ): loc, _, _ = unwrap_write_loc(loc_info) maybe_detect_oob(loc, 0, self.size + self.page_size, "set_kv_buffer (MLA)") - layer_id = layer.layer_id + layer_id = ( + layer_id_override if layer_id_override is not None else layer.layer_id + ) assert not self.dsa_kv_cache_store_fp8 parallel = get_parallel() if parallel.dcp_enabled: @@ -4201,6 +4204,7 @@ class MLATokenToKVPool(KVCache): loc: torch.Tensor, cache_k_nope: torch.Tensor, cache_k_rope: torch.Tensor, + layer_id_override: Optional[int] = None, ): # loc is widened under DCP; the kernel divides by the world size itself. maybe_detect_oob( @@ -4209,7 +4213,9 @@ class MLATokenToKVPool(KVCache): (self.size + self.page_size) * get_parallel().attn_dcp_size, "set_mla_kv_buffer (MLA)", ) - layer_id = layer.layer_id + layer_id = ( + layer_id_override if layer_id_override is not None else layer.layer_id + ) self._write_mla_kv_buffer( self.kv_buffer[layer_id - self.start_layer], loc, diff --git a/python/sglang/srt/mem_cache/swa_memory_pool.py b/python/sglang/srt/mem_cache/swa_memory_pool.py index 8b7cf4821..e934baf6d 100644 --- a/python/sglang/srt/mem_cache/swa_memory_pool.py +++ b/python/sglang/srt/mem_cache/swa_memory_pool.py @@ -6,8 +6,10 @@ import torch from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.mem_cache.base_swa_memory_pool import BaseSWAKVPool from sglang.srt.mem_cache.memory_pool import ( + DSATokenToKVPool, KVCache, MHATokenToKVPool, + MLATokenToKVPool, unwrap_write_loc, ) from sglang.srt.mem_cache.utils import maybe_init_custom_mem_pool @@ -17,7 +19,12 @@ GB = 1024 * 1024 * 1024 class SWAKVPool(BaseSWAKVPool): - """KV cache with separate pools for full and SWA attention layers.""" + """Hybrid full/SWA cache composed from independently configurable pools. + + The default remains two MHA pools. Supplying ``full_kv_pool_class`` and + ``swa_kv_pool_class`` enables other KV cache families, including MLA/DSA, + without adding model-specific behavior to the pool selector. + """ def __init__( self, @@ -31,6 +38,10 @@ class SWAKVPool(BaseSWAKVPool): full_attention_layer_ids: List[int], device: str, token_to_kv_pool_class: KVCache = MHATokenToKVPool, + full_kv_pool_class: Optional[type] = None, + swa_kv_pool_class: Optional[type] = None, + full_kv_pool_kwargs: Optional[dict] = None, + swa_kv_pool_kwargs: Optional[dict] = None, **kwargs, ): self.size = size @@ -46,35 +57,58 @@ class SWAKVPool(BaseSWAKVPool): self.page_size = page_size self.layer_transfer_counter = None - kwargs["page_size"] = page_size - kwargs["enable_memory_saver"] = False - kwargs["head_num"] = head_num - kwargs["head_dim"] = head_dim - kwargs["device"] = device - # for disagg with nvlink self.enable_custom_mem_pool, self.custom_mem_pool, _ = ( maybe_init_custom_mem_pool(device=self.device) ) - full_pool_kwargs = kwargs.copy() - full_pool_kwargs.pop("swa_head_num", None) - full_pool_kwargs.pop("swa_head_dim", None) - full_pool_kwargs.pop("swa_v_head_dim", None) - self.full_kv_pool = token_to_kv_pool_class( + full_kv_pool_class = full_kv_pool_class or token_to_kv_pool_class + swa_kv_pool_class = swa_kv_pool_class or token_to_kv_pool_class + common_kwargs = { + "page_size": page_size, + "enable_memory_saver": False, + "device": device, + } + if full_kv_pool_kwargs is None: + full_kv_pool_kwargs = { + **common_kwargs, + "head_num": head_num, + "head_dim": head_dim, + "allocation_label": "Full", + **kwargs, + } + full_kv_pool_kwargs.pop("swa_head_num", None) + full_kv_pool_kwargs.pop("swa_head_dim", None) + full_kv_pool_kwargs.pop("swa_v_head_dim", None) + if swa_kv_pool_kwargs is None: + swa_kv_pool_kwargs = { + **common_kwargs, + "head_num": head_num, + "head_dim": head_dim, + "allocation_label": "SWA", + **kwargs, + } + + self.full_kv_pool = full_kv_pool_class( size=size, dtype=dtype, layer_num=self.full_layer_nums, - allocation_label="Full", - **full_pool_kwargs, + **full_kv_pool_kwargs, ) - self.swa_kv_pool = token_to_kv_pool_class( + self.swa_kv_pool = swa_kv_pool_class( size=size_swa, dtype=dtype, layer_num=self.swa_layer_nums, - allocation_label="SWA", - **kwargs, + **swa_kv_pool_kwargs, ) + self.dsa_kv_cache_store_fp8 = False + self.kv_cache_dim = None + self.index_head_dim = None + if isinstance(self.full_kv_pool, MLATokenToKVPool): + self.dsa_kv_cache_store_fp8 = self.full_kv_pool.dsa_kv_cache_store_fp8 + self.kv_cache_dim = self.full_kv_pool.kv_cache_dim + if isinstance(self.full_kv_pool, DSATokenToKVPool): + self.index_head_dim = self.full_kv_pool.index_head_dim # {layer_id: (index, is_swa_layer)} self.layers_mapping: Dict[int, Tuple[int, bool]] = {} for full_attn_layer_id, global_layer_id in enumerate(full_attention_layer_ids): @@ -123,10 +157,19 @@ class SWAKVPool(BaseSWAKVPool): self.layer_transfer_counter.wait_until(layer_id - self.start_layer) def get_kv_size_bytes(self): - k_size, v_size = self.full_kv_pool.get_kv_size_bytes() - k_size_swa, v_size_swa = self.swa_kv_pool.get_kv_size_bytes() + def split_size(pool): + size = pool.get_kv_size_bytes() + return size if isinstance(size, tuple) else (size, 0) + + k_size, v_size = split_size(self.full_kv_pool) + k_size_swa, v_size_swa = split_size(self.swa_kv_pool) return k_size + k_size_swa, v_size + v_size_swa + def is_mla(self) -> bool: + return isinstance(self.full_kv_pool, MLATokenToKVPool) and isinstance( + self.swa_kv_pool, MLATokenToKVPool + ) + def get_contiguous_buf_infos(self): full_kv_data_ptrs, full_kv_data_lens, full_kv_item_lens = ( self.full_kv_pool.get_contiguous_buf_infos() @@ -200,21 +243,22 @@ class SWAKVPool(BaseSWAKVPool): loc, swa_loc, _ = unwrap_write_loc(loc_info) layer_id = layer.layer_id layer_id_pool, is_swa_layer = self.layers_mapping[layer_id] + pool = self.swa_kv_pool if is_swa_layer else self.full_kv_pool if is_swa_layer: # swa_loc is the full->SWA translation, computed once per forward by # the attention backend; set_kv_buffer never translates internally. assert swa_loc is not None - self.swa_kv_pool.set_kv_buffer( + loc = swa_loc + if isinstance(pool, MLATokenToKVPool): + pool.set_kv_buffer( None, - swa_loc, + loc, cache_k, cache_v, - k_scale, - v_scale, layer_id_override=layer_id_pool, ) else: - self.full_kv_pool.set_kv_buffer( + pool.set_kv_buffer( None, loc, cache_k, @@ -224,6 +268,60 @@ class SWAKVPool(BaseSWAKVPool): layer_id_override=layer_id_pool, ) + def set_mla_kv_buffer( + self, + layer: RadixAttention, + loc_info, + cache_k_nope: torch.Tensor, + cache_k_rope: torch.Tensor, + ): + loc, swa_loc, _ = unwrap_write_loc(loc_info) + layer_id_pool, is_swa_layer = self.layers_mapping[layer.layer_id] + pool = self.swa_kv_pool if is_swa_layer else self.full_kv_pool + if is_swa_layer: + assert swa_loc is not None + loc = swa_loc + if not isinstance(pool, MLATokenToKVPool): + raise TypeError(f"Layer {layer.layer_id} is not backed by an MLA KV pool") + pool.set_mla_kv_buffer( + None, + loc, + cache_k_nope, + cache_k_rope, + layer_id_override=layer_id_pool, + ) + + def get_index_k_with_scale_buffer(self, layer_id: int) -> torch.Tensor: + layer_id_pool, is_swa_layer = self.layers_mapping[layer_id] + assert not is_swa_layer + return self.full_kv_pool.get_index_k_with_scale_buffer(layer_id_pool) + + def get_index_k_continuous(self, layer_id: int, *args, **kwargs): + layer_id_pool, is_swa_layer = self.layers_mapping[layer_id] + assert not is_swa_layer + return self.full_kv_pool.get_index_k_continuous(layer_id_pool, *args, **kwargs) + + def get_index_k_scale_continuous(self, layer_id: int, *args, **kwargs): + layer_id_pool, is_swa_layer = self.layers_mapping[layer_id] + assert not is_swa_layer + return self.full_kv_pool.get_index_k_scale_continuous( + layer_id_pool, *args, **kwargs + ) + + def get_index_k_scale_buffer(self, layer_id: int, *args, **kwargs): + layer_id_pool, is_swa_layer = self.layers_mapping[layer_id] + assert not is_swa_layer + return self.full_kv_pool.get_index_k_scale_buffer( + layer_id_pool, *args, **kwargs + ) + + def set_index_k_scale_buffer(self, layer_id: int, *args, **kwargs): + layer_id_pool, is_swa_layer = self.layers_mapping[layer_id] + assert not is_swa_layer + return self.full_kv_pool.set_index_k_scale_buffer( + layer_id_pool, *args, **kwargs + ) + def move_kv_cache(self, tgt_loc: torch.Tensor, src_loc: torch.Tensor): self.full_kv_pool.move_kv_cache(tgt_loc, src_loc) tgt_loc_swa = self.translate_loc_from_full_to_swa(tgt_loc) diff --git a/python/sglang/srt/model_executor/forward_batch_deepseek_mha_mixin.py b/python/sglang/srt/model_executor/forward_batch_deepseek_mha_mixin.py index 5125cc7f0..4ee470352 100644 --- a/python/sglang/srt/model_executor/forward_batch_deepseek_mha_mixin.py +++ b/python/sglang/srt/model_executor/forward_batch_deepseek_mha_mixin.py @@ -128,11 +128,19 @@ class ForwardBatchDeepSeekMHAMixin: HybridLinearKVPool, MLATokenToKVPool, ) + from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool token_to_kv_pool = get_token_to_kv_pool() - assert isinstance(token_to_kv_pool, MLATokenToKVPool) or ( - isinstance(token_to_kv_pool, HybridLinearKVPool) - and isinstance(token_to_kv_pool.full_kv_pool, MLATokenToKVPool) + assert ( + isinstance(token_to_kv_pool, MLATokenToKVPool) + or ( + isinstance(token_to_kv_pool, HybridLinearKVPool) + and isinstance(token_to_kv_pool.full_kv_pool, MLATokenToKVPool) + ) + or ( + isinstance(token_to_kv_pool, SWAKVPool) + and isinstance(token_to_kv_pool.full_kv_pool, MLATokenToKVPool) + ) ), "Currently chunked prefix cache can only be used by Deepseek models" if not any(self.extend_prefix_lens_cpu): diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index a93c5d3d4..77fbbd1d7 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -514,7 +514,6 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): # Has to be None when cuda graph is captured. global_num_tokens_for_logprob_cpu: Optional[List[int]] = None global_num_tokens_for_logprob_gpu: Optional[torch.Tensor] = None - # For padding num_token_non_padded: Optional[torch.Tensor] = None # scalar tensor num_token_non_padded_cpu: int = None @@ -1535,9 +1534,35 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): dim=1, ) - # TODO: check if we need to pad other tensors + # Draft-extend padding uses the fixed per-request token width. + dummy_extend_len = 0 + if ( + self.spec_info is not None + and self.forward_mode.is_draft_extend_v2() + and self.spec_info.num_tokens_per_req > 0 + ): + dummy_extend_len = self.spec_info.num_tokens_per_req + if self.extend_seq_lens is not None: - self.extend_seq_lens = self._pad_tensor_to_size(self.extend_seq_lens, bs) + self.extend_seq_lens = self._pad_tensor_to_size( + self.extend_seq_lens, bs, value=dummy_extend_len + ) + if self.extend_prefix_lens is not None: + self.extend_prefix_lens = self._pad_tensor_to_size( + self.extend_prefix_lens, bs + ) + if self.extend_seq_lens_cpu is not None: + self.extend_seq_lens_cpu.extend( + [dummy_extend_len] * (bs - len(self.extend_seq_lens_cpu)) + ) + if self.extend_prefix_lens_cpu is not None: + self.extend_prefix_lens_cpu.extend( + [0] * (bs - len(self.extend_prefix_lens_cpu)) + ) + if self.extend_logprob_start_lens_cpu is not None: + self.extend_logprob_start_lens_cpu.extend( + [0] * (bs - len(self.extend_logprob_start_lens_cpu)) + ) if self.rids_int is not None: self.rids_int = self._pad_tensor_to_size(self.rids_int, bs) diff --git a/python/sglang/srt/model_executor/model_runner_components/spec_aux_hidden_state.py b/python/sglang/srt/model_executor/model_runner_components/spec_aux_hidden_state.py index c3513fd72..ce8940055 100644 --- a/python/sglang/srt/model_executor/model_runner_components/spec_aux_hidden_state.py +++ b/python/sglang/srt/model_executor/model_runner_components/spec_aux_hidden_state.py @@ -31,6 +31,8 @@ def _map_muse_target_layer_ids(*, target_hf_config, draft_hf_config, layer_ids): class SpecAuxHiddenStateConfig(msgspec.Struct, kw_only=True): eagle_use_aux_hidden_state: bool = False eagle_draft_num_layers: Optional[int] = None + # Draft layers whose KV cache uses the target SWA pool capacity. + eagle_draft_swa_num_layers: Optional[int] = None eagle_aux_hidden_state_layer_ids: Any = None dflash_use_aux_hidden_state: bool = False dflash_draft_num_layers: Optional[int] = None @@ -93,6 +95,14 @@ def _resolve_eagle_aux_hidden_state( ) ) + if ( + draft_model_config.is_hybrid_swa + and not draft_model_config.is_deepseek_v4_arch + ): + config.eagle_draft_swa_num_layers = len( + draft_model_config.swa_attention_layer_ids + ) + if spec_algorithm.is_eagle3(): config.eagle_use_aux_hidden_state = True try: diff --git a/python/sglang/srt/model_executor/pool_configurator.py b/python/sglang/srt/model_executor/pool_configurator.py index 8698cef49..136e725ef 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/python/sglang/srt/model_executor/pool_configurator.py @@ -21,6 +21,7 @@ import torch from sglang.srt.configs.hybrid_arch import mambaish_config from sglang.srt.configs.model_config import ( + AttentionArch, dsa_layer_skips_topk, get_dsa_index_head_dim, get_minimax_sparse_attention_config, @@ -431,7 +432,7 @@ class DefaultPoolConfigurator(MemoryPoolConfigurator): class HybridSWAPoolConfigurator(MemoryPoolConfigurator): - """Configurator for hybrid sliding window attention models (Gemma2, Command-R, MiMo). + """Configurator for MHA or MLA models with sliding-window layers. Splits available memory between full attention and SWA pools. Does NOT inherit DefaultPoolConfigurator — different coeff model. @@ -454,19 +455,46 @@ class HybridSWAPoolConfigurator(MemoryPoolConfigurator): self._sliding_window_size = kvc.sliding_window_size self._page_size = kvc.page_size - # Full layer per-token memory (bytes) - self._full_per_token = ( - model_config.get_num_kv_heads(tp_size) - * (model_config.head_dim + model_config.v_head_dim) - * kv_size - ) + if model_config.attention_arch == AttentionArch.MLA: + # MLA pool sizing uses latent dimensions rather than MHA heads. + from sglang.srt.mem_cache.kv_cache_configurator import ( + calculate_mla_kv_cache_dim, + ) - # SWA layer per-token memory (bytes) - self._swa_per_token = ( - model_config.get_swa_num_kv_heads(tp_size) - * (model_config.swa_head_dim + model_config.swa_v_head_dim) - * kv_size - ) + self._full_per_token = ( + calculate_mla_kv_cache_dim( + model_config=model_config, + kv_cache_dtype=kv_cache_dtype, + server_args=kvc.server_args, + ) + * kv_size + ) + if is_deepseek_dsa(model_config.hf_config): + index_head_dim = get_dsa_index_head_dim(model_config.hf_config) + index_elements = ( + index_head_dim + + index_head_dim // DSATokenToKVPool.quant_block_size * 4 + ) + self._full_per_token += index_elements * torch._utils._element_size( + DSATokenToKVPool.index_k_with_scale_buffer_dtype + ) + self._swa_per_token = ( + model_config.swa_kv_lora_rank + model_config.swa_qk_rope_head_dim + ) * kv_size + else: + # Full layer per-token memory (bytes) + self._full_per_token = ( + model_config.get_num_kv_heads(tp_size) + * (model_config.head_dim + model_config.v_head_dim) + * kv_size + ) + + # SWA layer per-token memory (bytes) + self._swa_per_token = ( + model_config.get_swa_num_kv_heads(tp_size) + * (model_config.swa_head_dim + model_config.swa_v_head_dim) + * kv_size + ) if self.kv_cache_dtype_str == "mxfp8": scale_block_size = 32 @@ -479,11 +507,9 @@ class HybridSWAPoolConfigurator(MemoryPoolConfigurator): * (model_config.swa_head_dim + model_config.swa_v_head_dim) ) // scale_block_size - # EAGLE/STANDALONE draft KV pool inherits max_total tokens with its - # full-attn layers; budget into the full term. A banded MTP depth - # (Inkling mtp_local_layer_ids) instead allocates an swa-geometry ring - # at FULL draft capacity, so budget those depths at swa_per_token. + # Draft KV tensors use full, SWA, or full-capacity SWA geometry. self._draft_full_layers_num = 0 + self._draft_swa_layers_num = 0 self._draft_swa_full_layers_num = 0 if ( kvc.spec_algorithm.is_eagle() or kvc.spec_algorithm.is_standalone() @@ -503,8 +529,18 @@ class HybridSWAPoolConfigurator(MemoryPoolConfigurator): if i < draft_layers ] ) - self._draft_swa_full_layers_num = banded_depths - self._draft_full_layers_num = draft_layers - banded_depths + self._draft_swa_full_layers_num = banded_depths + else: + draft_swa_layers = kvc.spec_aux_config.eagle_draft_swa_num_layers + if draft_swa_layers is not None: + self._draft_swa_layers_num = min( + max(int(draft_swa_layers), 0), draft_layers + ) + self._draft_full_layers_num = ( + draft_layers + - self._draft_swa_layers_num + - self._draft_swa_full_layers_num + ) self._draft_cell_size = _dflash_draft_cell_size(kvc) @@ -521,6 +557,7 @@ class HybridSWAPoolConfigurator(MemoryPoolConfigurator): self._cell_size = ( self._swa_per_token * self._swa_layers_num + self._full_per_token * self._draft_full_layers_num + + self._swa_per_token * self._draft_swa_layers_num + self._swa_per_token * self._draft_swa_full_layers_num + self._draft_cell_size ) @@ -531,7 +568,7 @@ class HybridSWAPoolConfigurator(MemoryPoolConfigurator): + self._swa_per_token * self._draft_swa_full_layers_num + self._swa_full_tokens_ratio * self._swa_per_token - * self._swa_layers_num + * (self._swa_layers_num + self._draft_swa_layers_num) + self._draft_cell_size ) @@ -667,7 +704,11 @@ class SWAChunkCapPoolConfigurator(HybridSWAPoolConfigurator): ) -> MemoryPoolConfig: # SWA pool sized tightly from the cap; the rest of the budget goes to full. swa_tokens = ceil_align(self._swa_cap, page_size) - fixed_swa_bytes = swa_tokens * self._swa_per_token * self._swa_layers_num + fixed_swa_bytes = ( + swa_tokens + * self._swa_per_token + * (self._swa_layers_num + self._draft_swa_layers_num) + ) full_cell_size = ( self._full_per_token * (self._full_layers_num + self._draft_full_layers_num) + self._swa_per_token * self._draft_swa_full_layers_num diff --git a/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py b/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py index 9899f9fb8..4d17d45e9 100644 --- a/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py +++ b/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py @@ -78,6 +78,14 @@ def _clone_if_runai_streamed_tensor(tensor: torch.Tensor) -> torch.Tensor: return tensor +def _get_indexer_weight_block_size( + quant_config: Optional[QuantizationConfig], +) -> List[int]: + if quant_config is not None and quant_config.weight_block_size is not None: + return quant_config.weight_block_size + return [128, 128] + + def _load_fused_indexer_wk( name: str, loaded_weight: torch.Tensor, @@ -99,8 +107,23 @@ def _load_fused_indexer_wk( return False if ".indexer.weights_proj." in name: - w = _clone_if_runai_streamed_tensor(loaded_weight) - fused_param.data[-w.shape[0] :].copy_(w) + is_scale = name.endswith(".weight_scale_inv") + if not is_scale and loaded_weight.dtype != torch.float8_e4m3fn: + w = _clone_if_runai_streamed_tensor(loaded_weight) + fused_param.data[-w.shape[0] :].copy_(w) + return True + + entry = pending.setdefault(fused_name + ".weights_proj", {}) + entry["scale" if is_scale else "weight"] = _clone_if_runai_streamed_tensor( + loaded_weight + ) + if "weight" in entry and "scale" in entry: + pending.pop(fused_name + ".weights_proj") + block_size = _get_indexer_weight_block_size(quant_config) + weights_bf16 = block_quant_dequant( + entry["weight"], entry["scale"], block_size, torch.bfloat16 + ) + fused_param.data[-weights_bf16.shape[0] :].copy_(weights_bf16) return True # wk: a bf16 checkpoint copies straight in; block-fp8 needs weight + scale. @@ -116,7 +139,7 @@ def _load_fused_indexer_wk( ) if "weight" in entry and "scale" in entry: pending.pop(fused_name) - block_size = getattr(quant_config, "weight_block_size", None) or [128, 128] + block_size = _get_indexer_weight_block_size(quant_config) wk_bf16 = block_quant_dequant( entry["weight"], entry["scale"], block_size, torch.bfloat16 ) @@ -547,8 +570,10 @@ class DeepseekV2WeightLoaderMixin: ) if selected_quant_config is None: selected_quant_config = self.quant_config - weight_block_size = getattr( - selected_quant_config, "weight_block_size", None + weight_block_size = ( + selected_quant_config.weight_block_size + if selected_quant_config is not None + else None ) if weight_block_size is not None: assert hasattr(self_attn.kv_b_proj, "weight_scale_inv") or hasattr( @@ -571,8 +596,10 @@ class DeepseekV2WeightLoaderMixin: # In multiple weight loading scenarios (e.g. RL), we need to inverse the scale of the weights after the requantization happened at the first loading. if ( should_deepgemm_weight_requant_ue8m0( - weight_block_size=getattr( - self.quant_config, "weight_block_size", None + weight_block_size=( + self.quant_config.weight_block_size + if self.quant_config is not None + else None ) ) and weight_scale.format_ue8m0 @@ -624,16 +651,19 @@ class DeepseekV2WeightLoaderMixin: self_attn.w_scale = scale if w.dtype == torch.int8: - if hasattr(self.quant_config, "weight_block_size"): + weight_block_size = ( + self.quant_config.weight_block_size + if self.quant_config is not None + else None + ) + if weight_block_size is not None: # block-wise int8 need it - weight_block_size = self.quant_config.weight_block_size - if weight_block_size is not None: - assert hasattr(self_attn.kv_b_proj, "weight_scale_inv") - weight = w - weight_scale = self_attn.kv_b_proj.weight_scale_inv - w = int8_block_dequant( - weight, weight_scale, weight_block_size - ).to(torch.bfloat16) + assert hasattr(self_attn.kv_b_proj, "weight_scale_inv") + weight = w + weight_scale = self_attn.kv_b_proj.weight_scale_inv + w = int8_block_dequant(weight, weight_scale, weight_block_size).to( + torch.bfloat16 + ) else: # channel-wise int8 need it w = w.to(torch.bfloat16) * self_attn.kv_b_proj.weight_scale.to( diff --git a/python/sglang/srt/models/dots3.py b/python/sglang/srt/models/dots3.py new file mode 100644 index 000000000..fa551577c --- /dev/null +++ b/python/sglang/srt/models/dots3.py @@ -0,0 +1,37 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2023-2024 SGLang Team + +"""Registry entry point for the Dots3 model.""" + +from sglang.srt.models.dots3_common.modeling import ( + Dots3AttentionMLA, + Dots3AttnForwardMethod, + Dots3DecoderLayer, + Dots3LanguageModelForCausalLM, + Dots3MLP, + Dots3Model, + Dots3MoE, + Dots3MoEGate, + Dots3NoteForCausalLM, + DotsNoteOmniForConditionalGeneration, + DotsNoteOmniThinkerForConditionalGeneration, + get_attention_sliding_window_size, +) + +EntryClass = [Dots3NoteForCausalLM] + +__all__ = [ + "Dots3AttentionMLA", + "Dots3AttnForwardMethod", + "Dots3DecoderLayer", + "Dots3LanguageModelForCausalLM", + "Dots3MLP", + "Dots3Model", + "Dots3MoE", + "Dots3MoEGate", + "Dots3NoteForCausalLM", + "DotsNoteOmniForConditionalGeneration", + "DotsNoteOmniThinkerForConditionalGeneration", + "EntryClass", + "get_attention_sliding_window_size", +] diff --git a/python/sglang/srt/models/dots3_common/__init__.py b/python/sglang/srt/models/dots3_common/__init__.py new file mode 100644 index 000000000..4343b4aba --- /dev/null +++ b/python/sglang/srt/models/dots3_common/__init__.py @@ -0,0 +1 @@ +"""Shared implementation modules for Dots3 models.""" diff --git a/python/sglang/srt/models/dots3_common/dots_omni_audio.py b/python/sglang/srt/models/dots3_common/dots_omni_audio.py new file mode 100644 index 000000000..c2c2d2c47 --- /dev/null +++ b/python/sglang/srt/models/dots3_common/dots_omni_audio.py @@ -0,0 +1,1027 @@ +"""Dots-path speech encoder for inference only (single GPU). + +Ported from cybertron_alm ``dots_audio_encoder/modeling_whisper.py``. +Upstream ``WhisperEncoder`` is exposed as :class:`DotsSpeechEncoder`. +""" + +import math +from functools import lru_cache +from typing import ClassVar + +import torch +import torch.nn.functional as F +from torch import nn +from transformers.activations import ACT2FN +from transformers.audio_utils import mel_filter_bank +from transformers.modeling_outputs import BaseModelOutput +from transformers.modeling_utils import PreTrainedModel +from transformers.models.whisper.configuration_whisper import WhisperConfig +from transformers.utils import logging + +from sglang.kernels.ops.attention.flash_attention import flash_attn_varlen_func +from sglang.srt.layers.rotary_embedding.utils import rotate_neox + +logger = logging.get_logger(__name__) + +__all__ = [ + "DotsSpeechEncoder", + "DotsWhisperConfig", + "OmniAudioConfig", + "OmniAudioModel", + "compute_audio_token_length", +] + + +class DotsWhisperConfig(WhisperConfig): + """Whisper configuration with the fields required by the Dots encoder.""" + + def __init__( + self, + *, + use_causal: bool = False, + use_rms_norm: bool = False, + use_latent_input: bool = False, + use_conv2d_stem: bool = False, + latent_dim: int | None = None, + downsample_hidden_size: int = 480, + use_rope: bool = False, + rope_parameters: dict | None = None, + conv_chunksize: int = 500, + conv_stem_gradient_checkpointing: bool = False, + conv_bucket_step: int | None = None, + conv_bucket_max_elements: int | None = None, + **kwargs, + ): + super().__init__(**kwargs) + self.use_causal = use_causal + self.use_rms_norm = use_rms_norm + self.use_latent_input = use_latent_input + self.use_conv2d_stem = use_conv2d_stem + self.latent_dim = latent_dim + self.downsample_hidden_size = downsample_hidden_size + self.use_rope = use_rope + self.rope_parameters = rope_parameters or {} + self.conv_chunksize = conv_chunksize + self.conv_stem_gradient_checkpointing = conv_stem_gradient_checkpointing + self.conv_bucket_step = conv_bucket_step + self.conv_bucket_max_elements = conv_bucket_max_elements + + +class RMSNorm(nn.Module): + def __init__(self, dim: int, eps: float = 1e-6): + super().__init__() + self.weight = nn.Parameter(torch.ones(dim)) + self.eps = eps + + def forward(self, x: torch.Tensor) -> torch.Tensor: + var = x.pow(2).mean(-1, keepdim=True) + x = x * torch.rsqrt(var + self.eps) + return self.weight * x + + +def swiglu(x: torch.Tensor) -> torch.Tensor: + x1, x2 = x.chunk(2, dim=-1) + return torch.nn.functional.silu(x1) * x2 + + +class RotaryEmbedding(nn.Module): + def __init__( + self, + head_dim: int, + rope_parameters: dict, + base_seq_len: int = 0, + ): + super().__init__() + self.partial_rotary_factor = float( + rope_parameters.get("partial_rotary_factor", 1.0) + ) + self.rope_theta = float(rope_parameters.get("rope_theta", 10000.0)) + self.rope_type = rope_parameters.get("rope_type", "default") + rotary_dim = int(head_dim * self.partial_rotary_factor) + # keep even rotary dimension to align cos/sin pairs + self.rotary_dim = (rotary_dim // 2) * 2 + self.attention_scaling = 1.0 + if self.rope_type != "default": + logger.warning( + f"RoPE type {self.rope_type} not implemented in WhisperLv3, fallback to default." + ) + if self.rotary_dim > 0: + inv_freq = 1.0 / ( + self.rope_theta + ** ( + torch.arange(0, self.rotary_dim, 2, dtype=torch.float) + / max(self.rotary_dim, 1) + ) + ) + else: + inv_freq = torch.tensor([]) + self.register_buffer("inv_freq", inv_freq, persistent=False) + self.base_seq_len = max(0, int(base_seq_len)) + self._cache: ( + tuple[int, torch.dtype, torch.device, torch.Tensor, torch.Tensor] | None + ) = None + + @torch.no_grad() + def get_cos_sin( + self, position_ids: torch.Tensor, dtype: torch.dtype, device: torch.device + ) -> tuple[torch.Tensor | None, torch.Tensor | None]: + if self.rotary_dim == 0: + return None, None + seq_len = position_ids.shape[-1] + if position_ids.shape[0] == 1 and self._cache is not None: + cached_seq_len, cached_dtype, cached_device, cached_cos, cached_sin = ( + self._cache + ) + if ( + cached_seq_len >= seq_len + and cached_dtype == dtype + and cached_device == device + ): + return cached_cos[:, :seq_len, :], cached_sin[:, :seq_len, :] + + if position_ids.shape[0] == 1: + cache_seq_len = max(seq_len, self.base_seq_len) + position_ids = torch.arange(cache_seq_len, device=device)[None, :] + # [1, D/2, 1] + if self.inv_freq.dtype != torch.float32: + inv_freq = 1.0 / ( + self.rope_theta + ** ( + torch.arange( + 0, self.rotary_dim, 2, dtype=torch.float, device=device + ) + / max(self.rotary_dim, 1) + ) + ) + else: + inv_freq = self.inv_freq.to(device=device) + inv_freq_expanded = inv_freq[None, :, None] + # [B, 1, T] + position_ids_expanded = position_ids[:, None, :].float() + # Force float32 since bfloat16 loses precision on long contexts + # See https://github.com/huggingface/transformers/pull/29285 + device_type = device.type if isinstance(device, torch.device) else "cpu" + with torch.autocast(device_type=device_type, enabled=False): + freqs = ( + inv_freq_expanded.float() @ position_ids_expanded.float() + ).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() * self.attention_scaling + sin = emb.sin() * self.attention_scaling + cos = cos.to(dtype) + sin = sin.to(dtype) + if position_ids.shape[0] == 1: + self._cache = (position_ids.shape[-1], dtype, device, cos, sin) + return cos[:, :seq_len, :], sin[:, :seq_len, :] + return cos, sin + + +def apply_rotary_pos_emb( + q: torch.Tensor, + k: torch.Tensor, + cos: torch.Tensor | None, + sin: torch.Tensor | None, +) -> tuple[torch.Tensor, torch.Tensor]: + if cos is None or sin is None: + return q, k + + if q.dim() not in (3, 4): + return q, k + + rotary_dim = cos.shape[-1] + unsqueeze_dim = 2 if q.dim() == 4 else 1 + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:] + k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:] + q_embed = (q_rot * cos) + (rotate_neox(q_rot) * sin) + k_embed = (k_rot * cos) + (rotate_neox(k_rot) * sin) + return torch.cat((q_embed, q_pass), dim=-1), torch.cat((k_embed, k_pass), dim=-1) + + +class WhisperAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__( + self, + embed_dim: int, + num_heads: int, + dropout: float = 0.0, + is_decoder: bool = False, + bias: bool = True, + ): + super().__init__() + self.embed_dim = embed_dim + self.num_heads = num_heads + self.dropout = dropout + self.head_dim = embed_dim // num_heads + + if (self.head_dim * num_heads) != self.embed_dim: + raise ValueError( + f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}" + f" and `num_heads`: {num_heads})." + ) + self.scaling = self.head_dim**-0.5 + self.is_decoder = is_decoder + + self.k_proj = nn.Linear(embed_dim, embed_dim, bias=False) + self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias) + self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias) + self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias) + + def forward_flash_attn( + self, + hidden_states, + cu_seqlens_q=None, + cu_seqlens_kv=None, + max_seqlen_q=None, + max_seqlen_kv=None, + output_attentions=False, + rotary_cos=None, + rotary_sin=None, + ): + """Dense eager attention with SGLang FA3 for packed variable-length input.""" + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + if cu_seqlens_q is None: + bsz, tgt_len, _ = hidden_states.size() + query_states = query_states.view( + bsz, tgt_len, self.num_heads, self.head_dim + ) + key_states = key_states.view(bsz, tgt_len, self.num_heads, self.head_dim) + value_states = value_states.view( + bsz, tgt_len, self.num_heads, self.head_dim + ) + if rotary_cos is not None and rotary_sin is not None: + cos = rotary_cos[:, :tgt_len, :] + sin = rotary_sin[:, :tgt_len, :] + query_states, key_states = apply_rotary_pos_emb( + query_states, key_states, cos, sin + ) + attn_output, attn_probs = self._eager_attention( + query_states, key_states, value_states, output_attentions + ) + attn_output = attn_output.view(bsz, tgt_len, self.embed_dim) + else: + query_states = query_states.view(-1, self.num_heads, self.head_dim) + key_states = key_states.view(-1, self.num_heads, self.head_dim) + value_states = value_states.view(-1, self.num_heads, self.head_dim) + if rotary_cos is not None and rotary_sin is not None: + query_states, key_states = apply_rotary_pos_emb( + query_states, key_states, rotary_cos, rotary_sin + ) + result = flash_attn_varlen_func( + query_states, + key_states, + value_states, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_kv, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_kv, + softmax_scale=self.scaling, + causal=self.is_decoder, + return_softmax_lse=output_attentions, + ) + if isinstance(result, tuple): + attn_output = result[0] + attn_probs = result[1] if output_attentions else None + else: + attn_output = result + attn_probs = None + attn_output = attn_output.view(-1, self.embed_dim) + + return self.out_proj(attn_output), attn_probs + + def _eager_attention(self, query, key, value, output_attentions): + query = query.transpose(1, 2) + key = key.transpose(1, 2) + value = value.transpose(1, 2) + scores = torch.matmul(query, key.transpose(-2, -1)) * self.scaling + if self.is_decoder: + seq_len = scores.shape[-1] + causal_mask = torch.ones( + seq_len, seq_len, dtype=torch.bool, device=scores.device + ).triu(1) + scores.masked_fill_(causal_mask, torch.finfo(scores.dtype).min) + probabilities = torch.softmax(scores, dim=-1, dtype=torch.float32).to( + query.dtype + ) + output = torch.matmul(probabilities, value).transpose(1, 2) + return output, probabilities if output_attentions else None + + +# Copied from transformers.models.mbart.modeling_mbart.MBartEncoderLayer with MBart->Whisper +class WhisperEncoderLayer(nn.Module): + def __init__(self, config: DotsWhisperConfig): + super().__init__() + self.embed_dim = config.d_model + use_causal = config.use_causal + self.self_attn = WhisperAttention( + embed_dim=self.embed_dim, + num_heads=config.encoder_attention_heads, + dropout=config.attention_dropout, + is_decoder=use_causal, # enable causal attention when use_causal=True + ) + self.use_causal = use_causal + norm_cls = RMSNorm if config.use_rms_norm else nn.LayerNorm + self.self_attn_layer_norm = norm_cls(self.embed_dim) + self.dropout = config.dropout + self.activation_fn = ( + ACT2FN[config.activation_function] + if config.activation_function != "swiglu" + else swiglu + ) + self.use_swiglu = config.activation_function == "swiglu" + self.activation_dropout = config.activation_dropout + ffn_dim = config.encoder_ffn_dim + fc1_out = ffn_dim * 2 if self.use_swiglu else ffn_dim + self.fc1 = nn.Linear(self.embed_dim, fc1_out) + self.fc2 = nn.Linear(ffn_dim, self.embed_dim) + self.final_layer_norm = norm_cls(self.embed_dim) + + def forward( + self, + hidden_states: torch.Tensor, + cu_seqlens_q: torch.Tensor = None, + cu_seqlens_kv: torch.Tensor = None, + max_seqlen_q: int | None = None, + max_seqlen_kv: int | None = None, + output_attentions: bool = False, + rotary_cos: torch.Tensor | None = None, + rotary_sin: torch.Tensor | None = None, + ) -> torch.Tensor: + """ + Args: + hidden_states (`torch.FloatTensor`): input to the layer of shape `(seq_len, batch, embed_dim)` + attention_mask (`torch.FloatTensor`): attention mask of size + `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. + layer_head_mask (`torch.FloatTensor`): mask for attention heads in a given layer of size + `(encoder_attention_heads,)`. + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more detail. + """ + residual = hidden_states + hidden_states = self.self_attn_layer_norm(hidden_states) + + hidden_states, attn_weights = self.self_attn.forward_flash_attn( + hidden_states=hidden_states, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_kv, + max_seqlen_q=max_seqlen_q, + max_seqlen_kv=max_seqlen_kv, + output_attentions=output_attentions, + rotary_cos=rotary_cos, + rotary_sin=rotary_sin, + ) + hidden_states = nn.functional.dropout( + hidden_states, p=self.dropout, training=self.training + ) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.final_layer_norm(hidden_states) + if self.use_swiglu: + hidden_states = swiglu(self.fc1(hidden_states)) + else: + hidden_states = self.activation_fn(self.fc1(hidden_states)) + hidden_states = nn.functional.dropout( + hidden_states, p=self.activation_dropout, training=self.training + ) + hidden_states = self.fc2(hidden_states) + hidden_states = nn.functional.dropout( + hidden_states, p=self.dropout, training=self.training + ) + hidden_states = residual + hidden_states + + if hidden_states.dtype == torch.float16 and ( + torch.isinf(hidden_states).any() or torch.isnan(hidden_states).any() + ): + clamp_value = torch.finfo(hidden_states.dtype).max - 1000 + hidden_states = torch.clamp( + hidden_states, min=-clamp_value, max=clamp_value + ) + + outputs = (hidden_states,) + + if output_attentions: + outputs += (attn_weights,) + + return outputs + + +class DotsSpeechPreTrainedModel(PreTrainedModel): + config_class = DotsWhisperConfig + base_model_prefix = "model" + main_input_name = "input_features" + supports_gradient_checkpointing = False + _no_split_modules: ClassVar[list[str]] = ["WhisperEncoderLayer"] + + def _init_weights(self, module): + std = self.config.init_std + if isinstance(module, (nn.Linear, nn.Conv1d, nn.Conv2d)): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + + +class DotsSpeechEncoder(DotsSpeechPreTrainedModel): + """ + Transformer encoder consisting of *config.encoder_layers* self attention layers. Each layer is a + [`WhisperEncoderLayer`]. + + Args: + config: WhisperConfig + """ + + def __init__(self, config: DotsWhisperConfig): + if not isinstance(config, DotsWhisperConfig): + raise TypeError( + "DotsSpeechEncoder requires DotsWhisperConfig, not a plain " + "transformers.WhisperConfig." + ) + super().__init__(config) + self.dropout = config.dropout + + embed_dim = config.d_model + self.use_latent_input = config.use_latent_input + self.use_causal = config.use_causal + self.use_conv2d_stem = config.use_conv2d_stem + latent_dim = config.latent_dim + if self.use_latent_input and latent_dim is None: + raise ValueError( + "DotsSpeechEncoder: use_latent_input=True requires config.latent_dim" + ) + if self.use_conv2d_stem and self.use_latent_input: + raise ValueError( + "DotsSpeechEncoder: use_conv2d_stem and use_latent_input are mutually exclusive" + ) + self.num_mel_bins = ( + latent_dim if self.use_latent_input and latent_dim is not None else 128 + ) + self.padding_idx = config.pad_token_id + self.max_source_positions = config.max_source_positions + self.embed_scale = math.sqrt(embed_dim) if config.scale_embedding else 1.0 + + if self.use_conv2d_stem: + # Conv2D stem: 3 layers of stride-2 for 8x downsampling + dhs = config.downsample_hidden_size + # Causal: keep freq padding, remove time padding (handled by pad(14,0) in forward) + conv_padding = (1, 0) if self.use_causal else 1 + self.conv2d1 = nn.Conv2d( + 1, dhs, kernel_size=3, stride=2, padding=conv_padding + ) + self.conv2d2 = nn.Conv2d( + dhs, dhs, kernel_size=3, stride=2, padding=conv_padding + ) + self.conv2d3 = nn.Conv2d( + dhs, dhs, kernel_size=3, stride=2, padding=conv_padding + ) + # After 3x stride-2 on freq=128: 128→64→32→16; linear projects dhs*16 → embed_dim + freq_after = self.num_mel_bins + for _ in range(3): + freq_after = (freq_after + 1) // 2 + self.conv_out = nn.Linear(dhs * freq_after, embed_dim, bias=False) + self.conv1 = None + self.conv2 = None + elif self.use_latent_input: + # Latent path keeps length (stride=1) and applies GLU after conv1 + self.conv1 = nn.Conv1d( + self.num_mel_bins, embed_dim * 2, kernel_size=3, stride=1, padding=1 + ) + self.conv2 = nn.Conv1d( + embed_dim, embed_dim, kernel_size=3, stride=1, padding=1 + ) + else: + self.conv1 = nn.Conv1d( + self.num_mel_bins, embed_dim, kernel_size=3, padding=1 + ) + self.conv2 = nn.Conv1d( + embed_dim, embed_dim, kernel_size=3, stride=2, padding=1 + ) + + self.use_rope = config.use_rope + rope_parameters = config.rope_parameters + self.rope_parameters = rope_parameters + if self.use_rope: + head_dim = embed_dim // config.encoder_attention_heads + self.rotary_embedding = RotaryEmbedding( + head_dim, + rope_parameters, + base_seq_len=self.max_source_positions, + ) + self.embed_positions = None + else: + self.embed_positions = nn.Embedding(self.max_source_positions, embed_dim) + + self.layers = nn.ModuleList( + [WhisperEncoderLayer(config) for _ in range(config.encoder_layers)] + ) + norm_cls = RMSNorm if config.use_rms_norm else nn.LayerNorm + self.layer_norm = norm_cls(config.d_model) + + self.post_init() + + def get_input_embeddings(self) -> nn.Module: + if self.use_conv2d_stem: + return self.conv2d1 + return self.conv1 + + def set_input_embeddings(self, value: nn.Module): + self.conv1 = value + + @staticmethod + def _temporal_mask(feat, valid_lens): + """Zero out temporal positions >= valid_lens. feat: [B,C,F,T], valid_lens: [B] tensor.""" + T = feat.shape[-1] + mask = torch.arange(T, device=feat.device)[None, :] < valid_lens[:, None] + return mask[:, None, None, :] + + def _conv2d_stem_one_chunk(self, chunk, chunk_valid_mel_lens=None): + """Run 3x Conv2d(stride=2) + GELU with per-layer masking.""" + # Causal: left-pad time by 14 = 2 + 4 + 8 (receptive field of the 3-layer + # stride-2 stack mapped back to input time). Combined with conv padding=(1,0), + # this makes the whole conv2d stem strictly causal. + if self.use_causal: + chunk = nn.functional.pad(chunk, (14, 0)) + if chunk_valid_mel_lens is not None: + chunk_valid_mel_lens = chunk_valid_mel_lens + 14 + # Step 0: mask mel — silence_mel(-1.5) → 0 + if chunk_valid_mel_lens is not None: + chunk = chunk * self._temporal_mask(chunk, chunk_valid_mel_lens) + chunk = nn.functional.gelu(self.conv2d1(chunk)) + # Step 1: mask conv1 output — gelu(conv(0)+bias) → 0 + if chunk_valid_mel_lens is not None: + chunk_valid_mel_lens = (chunk_valid_mel_lens + 1) // 2 + chunk = chunk * self._temporal_mask(chunk, chunk_valid_mel_lens) + chunk = nn.functional.gelu(self.conv2d2(chunk)) + # Step 2: mask conv2 output — gelu(bias) → 0 + if chunk_valid_mel_lens is not None: + chunk_valid_mel_lens = (chunk_valid_mel_lens + 1) // 2 + chunk = chunk * self._temporal_mask(chunk, chunk_valid_mel_lens) + chunk = nn.functional.gelu(self.conv2d3(chunk)) + # Step 3: mask conv3 output — gelu(bias) → 0 + if chunk_valid_mel_lens is not None: + chunk_valid_mel_lens = (chunk_valid_mel_lens + 1) // 2 + chunk = chunk * self._temporal_mask(chunk, chunk_valid_mel_lens) + return chunk + + def _forward_conv2d_stem( + self, input_features, input_seq_lens=None, audio_sample_lens=None + ): + """Conv2D stem: 3x Conv2d(stride=2) → 8x downsample. Returns [B, T/8, embed_dim].""" + x = input_features.unsqueeze(1) # [B, 1, 128, T] + + if audio_sample_lens is not None: + hop_length = 160 + if not isinstance(audio_sample_lens, torch.Tensor): + audio_sample_lens = torch.tensor(audio_sample_lens, device=x.device) + valid_mel_lens = audio_sample_lens.to(x.device) // hop_length + else: + valid_mel_lens = None + x = self._conv2d_stem_one_chunk(x, valid_mel_lens) + # [B, dhs, F_out, T/8] → [B, T/8, dhs*freq_after] → [B, T/8, embed_dim] + batch, channels, frequency, time = x.shape + x = x.permute(0, 3, 1, 2).reshape(batch, time, channels * frequency) + inputs_embeds = self.conv_out(x) + + return inputs_embeds + + def _forward_conv1d_stem(self, input_features): + """Standard Conv1d stem: conv1 + conv2(stride=2), 2x downsample. Returns [B, T/2, embed_dim].""" + if self.use_causal: + x = nn.functional.pad(input_features, (2, 0)) + x = nn.functional.gelu(self.conv1(x)) + x = nn.functional.pad(x, (2, 0)) + x = nn.functional.gelu(self.conv2(x)) + else: + x = nn.functional.gelu(self.conv1(input_features)) + x = nn.functional.gelu(self.conv2(x)) + return x.permute(0, 2, 1) + + def _forward_latent_stem(self, input_features): + """Latent input stem: conv1+GLU + conv2, stride=1, no downsample. Returns [B, T, embed_dim].""" + if self.use_causal: + x = nn.functional.pad(input_features, (2, 0)) + x = nn.functional.glu(self.conv1(x), dim=1) + x = nn.functional.pad(x, (2, 0)) + x = nn.functional.gelu(self.conv2(x)) + else: + x = nn.functional.glu(self.conv1(input_features), dim=1) + x = nn.functional.gelu(self.conv2(x)) + return x.permute(0, 2, 1) + + def forward( + self, + input_features, + input_seq_lens=None, + audio_sample_lens=None, + output_attentions=None, + output_hidden_states=None, + return_dict=None, + ): + """Mel `[B, n_mels, T]` → encoder hidden states. Inference-only.""" + output_attentions = ( + output_attentions + if output_attentions is not None + else self.config.output_attentions + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + if self.use_conv2d_stem: + inputs_embeds = self._forward_conv2d_stem( + input_features, input_seq_lens, audio_sample_lens + ) + elif self.use_latent_input: + inputs_embeds = self._forward_latent_stem(input_features) + else: + inputs_embeds = self._forward_conv1d_stem(input_features) + rotary_cos = None + rotary_sin = None + position_ids = torch.arange( + inputs_embeds.shape[1], device=inputs_embeds.device + )[None, :] + if self.use_rope: + rotary_cos, rotary_sin = self.rotary_embedding.get_cos_sin( + position_ids, inputs_embeds.dtype, inputs_embeds.device + ) + hidden_states = inputs_embeds + else: + embed_pos = self.embed_positions.weight[: inputs_embeds.shape[1]] + hidden_states = inputs_embeds + embed_pos + + encoder_states = () if output_hidden_states else None + all_attentions = () if output_attentions else None + + if input_seq_lens is not None: + # Build varlen metadata and pack valid tokens without per-sample cat loops. + # Read max on whatever device the caller provided: when it is a CPU + # tensor this avoids a device->host sync (one per forward). + max_seqlen_q = int(input_seq_lens.max().item()) + input_seq_lens = input_seq_lens.to( + device=hidden_states.device, dtype=torch.long + ) + B, S, D = hidden_states.shape + max_seqlen_kv = max_seqlen_q + cu_seqlens_q = torch.nn.functional.pad( + input_seq_lens.cumsum(0, dtype=torch.int32), (1, 0) + ) + cu_seqlens_kv = cu_seqlens_q + + token_positions = torch.arange(S, device=hidden_states.device)[None, :] + valid_token_mask = token_positions < input_seq_lens[:, None] + hidden_states = hidden_states[valid_token_mask] + if rotary_cos is not None and rotary_sin is not None: + packed_positions = token_positions.expand(B, S)[valid_token_mask] + rotary_cos = rotary_cos.squeeze(0).index_select(0, packed_positions) + rotary_sin = rotary_sin.squeeze(0).index_select(0, packed_positions) + else: + cu_seqlens_q = None + cu_seqlens_kv = None + max_seqlen_q = None + max_seqlen_kv = None + + for encoder_layer in self.layers: + if output_hidden_states: + encoder_states = encoder_states + (hidden_states,) + layer_outputs = encoder_layer( + hidden_states, + cu_seqlens_q, + cu_seqlens_kv, + max_seqlen_q, + max_seqlen_kv, + output_attentions=output_attentions, + rotary_cos=rotary_cos, + rotary_sin=rotary_sin, + ) + hidden_states = layer_outputs[0] + if output_attentions: + all_attentions = all_attentions + (layer_outputs[1],) + + if input_seq_lens is not None: + # Recover packed varlen output directly on device. + recover_positions = torch.arange(max_seqlen_q, device=hidden_states.device)[ + None, : + ] + recover_mask = recover_positions < input_seq_lens[:, None] + recovered = hidden_states.new_zeros(B, max_seqlen_q, D) + recovered[recover_mask] = hidden_states + hidden_states = recovered + + hidden_states = self.layer_norm(hidden_states) + if output_hidden_states: + encoder_states = encoder_states + (hidden_states,) + + if not return_dict: + return tuple( + v + for v in [hidden_states, encoder_states, all_attentions] + if v is not None + ) + return BaseModelOutput( + last_hidden_state=hidden_states, + hidden_states=encoder_states, + attentions=all_attentions, + ) + + +SAMPLE_RATE = 16000 +N_FFT = 400 +HOP_LENGTH = 160 +DEFAULT_CHUNK_LENGTH_S = 60 +DEFAULT_CONV_TEMPORAL_STRIDE = 8 +DEFAULT_MERGE_FACTOR = 1 +N_SAMPLES = DEFAULT_CHUNK_LENGTH_S * SAMPLE_RATE + + +class OmniAudioConfig: + def __init__(self, **kwargs): + self.encoder_type = kwargs.get("encoder_type", "dots") + self.whisper_config = kwargs.get("whisper_config", {}) + self.whisper_adapter_in_dim = kwargs.get( + "whisper_adapter_in_dim", kwargs.get("adapter_in_dim", 1280) + ) + self.whisper_adapter_out_dim = kwargs.get( + "whisper_adapter_out_dim", kwargs.get("adapter_out_dim", 2048) + ) + self.sampling_rate = kwargs.get("sampling_rate", SAMPLE_RATE) + self.audio_comp_start = kwargs.get("audio_comp_start", "<|audio_comp_start|>") + self.audio_comp_span = kwargs.get("audio_comp_span", "<|audio_comp_pad|>") + self.audio_comp_end = kwargs.get("audio_comp_end", "<|audio_comp_end|>") + self.merge_factor = kwargs.get("merge_factor", DEFAULT_MERGE_FACTOR) + self.chunk_seconds = kwargs.get("chunk_seconds", DEFAULT_CHUNK_LENGTH_S) + self.use_conv2d_stem = kwargs.get("use_conv2d_stem", True) + self.use_latent_input = kwargs.get("use_latent_input", False) + self.latent_dim = kwargs.get("latent_dim") + self.use_rope = kwargs.get("use_rope", True) + self.use_rms_norm = kwargs.get("use_rms_norm", True) + self.use_causal = kwargs.get("use_causal", False) + self.downsample_hidden_size = kwargs.get("downsample_hidden_size", 480) + self.conv_chunksize = kwargs.get("conv_chunksize", 500) + self.conv_stem_gradient_checkpointing = kwargs.get( + "conv_stem_gradient_checkpointing", False + ) + self.conv_bucket_step = kwargs.get("conv_bucket_step", None) + self.conv_bucket_max_elements = kwargs.get("conv_bucket_max_elements", None) + self.rope_parameters = kwargs.get( + "rope_parameters", + { + "partial_rotary_factor": 0.5, + "rope_theta": 10000.0, + "rope_type": "default", + }, + ) + + @property + def conv_temporal_stride(self) -> int: + return 8 if self.use_conv2d_stem else 2 + + @property + def token_stride(self) -> int: + return HOP_LENGTH * self.conv_temporal_stride * self.merge_factor + + @property + def chunk_samples(self) -> int: + return int(self.chunk_seconds * SAMPLE_RATE) + + @property + def chunk_mel_frames(self) -> int: + return int(self.chunk_seconds * 100) + + +def pad_or_trim(array, length=N_SAMPLES, axis=-1): + if array.shape[axis] > length: + array = array.index_select( + dim=axis, index=torch.arange(length, device=array.device) + ) + if array.shape[axis] < length: + pad_widths = [(0, 0)] * array.ndim + pad_widths[axis] = (0, length - array.shape[axis]) + array = F.pad(array, [pad for sizes in pad_widths[::-1] for pad in sizes]) + return array + + +@lru_cache(maxsize=4) +def _mel_filters(device, n_mels=128): + filters = mel_filter_bank( + num_frequency_bins=1 + N_FFT // 2, + num_mel_filters=n_mels, + min_frequency=0.0, + max_frequency=float(SAMPLE_RATE) / 2.0, + sampling_rate=SAMPLE_RATE, + norm="slaney", + mel_scale="slaney", + ) + return torch.from_numpy(filters).T.contiguous().float().to(device) + + +@lru_cache(maxsize=4) +def _hann_window(device): + # Generate on CPU (default) then move, to preserve the exact reference + # window values. Direct on-device generation changes numerics slightly. + return torch.hann_window(N_FFT).to(device) + + +def log_mel_spectrogram(audio, n_mels=128): + window = _hann_window(audio.device) + stft = torch.stft(audio, N_FFT, HOP_LENGTH, window=window, return_complex=True) + magnitudes = stft[..., :-1].abs() ** 2 + filters = _mel_filters(audio.device, n_mels) + mel_spec = filters @ magnitudes + log_spec = torch.clamp(mel_spec, min=1e-10).log10() + log_spec = torch.maximum(log_spec, log_spec.max() - 8.0) + log_spec = (log_spec + 4.0) / 4.0 + return log_spec + + +def compute_audio_token_length( + num_samples, + *, + sample_rate: int = SAMPLE_RATE, + chunk_seconds: int = DEFAULT_CHUNK_LENGTH_S, + hop_length: int = HOP_LENGTH, + conv_temporal_stride: int = DEFAULT_CONV_TEMPORAL_STRIDE, + merge_factor: int = DEFAULT_MERGE_FACTOR, +): + stride = hop_length * conv_temporal_stride * merge_factor + total = 0 + time_step = 0 + while time_step * sample_rate < num_samples: + chunk_len = min( + num_samples - time_step * sample_rate, chunk_seconds * sample_rate + ) + total += math.ceil(chunk_len / stride) + time_step += chunk_seconds + return total + + +class DotsEncoderWithMask(nn.Module): + def __init__(self, config: OmniAudioConfig): + super().__init__() + whisper_config_kwargs = dict(config.whisper_config) + whisper_config_kwargs.update( + use_rope=config.use_rope, + rope_parameters=config.rope_parameters, + use_rms_norm=config.use_rms_norm, + use_causal=config.use_causal, + use_conv2d_stem=config.use_conv2d_stem, + use_latent_input=config.use_latent_input, + latent_dim=config.latent_dim, + downsample_hidden_size=config.downsample_hidden_size, + conv_chunksize=config.conv_chunksize, + conv_stem_gradient_checkpointing=(config.conv_stem_gradient_checkpointing), + conv_bucket_step=config.conv_bucket_step, + conv_bucket_max_elements=config.conv_bucket_max_elements, + ) + whisper_config = DotsWhisperConfig(**whisper_config_kwargs) + + self.speech_encoder = DotsSpeechEncoder(whisper_config) + self.merge_factor = config.merge_factor + self.chunk_seconds = config.chunk_seconds + self.chunk_samples = config.chunk_samples + self.chunk_mel_frames = config.chunk_mel_frames + self.conv_temporal_stride = config.conv_temporal_stride + + @property + def device(self): + return next(self.speech_encoder.parameters()).device + + def _forward_speech_encoder( + self, + mel_features: torch.Tensor, + input_seq_lens: torch.Tensor, + audio_sample_lens: list[int], + ) -> torch.Tensor: + """Run the eager speech encoder without server-side slicing/batching.""" + mel_features = mel_features.to(dtype=torch.bfloat16, device=self.device) + return self.speech_encoder( + mel_features, + return_dict=True, + input_seq_lens=input_seq_lens, + audio_sample_lens=audio_sample_lens, + ).last_hidden_state + + def encode_waveform(self, audio_waveform: torch.Tensor) -> torch.Tensor: + segments = [] + time_step = 0 + while time_step * SAMPLE_RATE < audio_waveform.shape[0]: + segments.append( + audio_waveform[ + time_step + * SAMPLE_RATE : (time_step + self.chunk_seconds) + * SAMPLE_RATE + ] + ) + time_step += self.chunk_seconds + + mel_features = [] + token_lens = [] + audio_sample_lens = [] + for audio_segment in segments: + segment_length = audio_segment.shape[0] + token_len = (segment_length - 1) // ( + HOP_LENGTH * self.conv_temporal_stride * self.merge_factor + ) + 1 + pad_audio = pad_or_trim(audio_segment.flatten(), length=self.chunk_samples) + mel = log_mel_spectrogram(pad_audio) + assert mel.shape[1] == self.chunk_mel_frames + mel_features.append(mel) + token_lens.append(token_len) + audio_sample_lens.append(segment_length) + + mel_features = torch.stack(mel_features, dim=0) + # Keep input_seq_lens on CPU: the conv2d bucket path reads it via + # ``.item()`` and CPU scalars avoid device->host syncs. The encoder's + # varlen path moves it to the device itself. + input_seq_lens = torch.tensor(token_lens, dtype=torch.long) * self.merge_factor + audio_embedding = self._forward_speech_encoder( + mel_features, input_seq_lens, audio_sample_lens + ) + + chunk_embeddings = [] + for idx, token_len in enumerate(token_lens): + chunk_embeddings.append( + audio_embedding[idx, : token_len * self.merge_factor, :] + ) + return torch.cat(chunk_embeddings, dim=0).unsqueeze(0) + + +class AudioAdapter(nn.Module): + def __init__(self, in_dim, out_dim): + super().__init__() + self.proj = nn.Sequential( + nn.LayerNorm(in_dim), + nn.Linear(in_dim, out_dim), + nn.GELU(), + nn.Linear(out_dim, out_dim), + ) + + def forward(self, x): + return self.proj(x) + + +class OmniAudioModel(nn.Module): + def __init__(self, config: OmniAudioConfig): + super().__init__() + if config.encoder_type != "dots": + raise ValueError("Dots omni only supports encoder_type='dots'") + self.merge_factor = config.merge_factor + self.audio_adapter = AudioAdapter( + config.whisper_adapter_in_dim, + config.whisper_adapter_out_dim, + ) + + self.dots_encoder = DotsEncoderWithMask(config) + + @property + def device(self): + return self.dots_encoder.device + + def _merge_embeddings(self, embedding: torch.Tensor) -> torch.Tensor: + if self.merge_factor <= 1: + return embedding + return embedding.reshape( + embedding.shape[0], + embedding.shape[1] // self.merge_factor, + embedding.shape[2] * self.merge_factor, + ) + + def _encode_single_audio(self, audio_waveform): + embedding = self.dots_encoder.encode_waveform(audio_waveform) + embedding = self._merge_embeddings(embedding) + embedding = self.audio_adapter(embedding) + return embedding.squeeze(0) + + def _split_waveforms(self, audio_inputs, lengths): + # Convert lengths to CPU ints once so the per-audio slicing below stays + # on the host and does not trigger a device->host sync per audio. + if isinstance(audio_inputs, list): + return audio_inputs + lengths_list = lengths.tolist() + waveforms = [] + audio_start = 0 + for length in lengths_list: + waveforms.append(audio_inputs[audio_start : audio_start + length]) + audio_start += length + return waveforms + + def forward(self, audio_inputs, lengths): + waveforms = self._split_waveforms(audio_inputs, lengths) + all_embeddings = [] + token_lengths = [] + for waveform in waveforms: + embedding = self._encode_single_audio(waveform) + all_embeddings.append(embedding) + token_lengths.append(embedding.shape[0]) + return torch.cat(all_embeddings, dim=0), token_lengths diff --git a/python/sglang/srt/models/dots3_common/dots_omni_towers.py b/python/sglang/srt/models/dots3_common/dots_omni_towers.py new file mode 100644 index 000000000..ea92ed489 --- /dev/null +++ b/python/sglang/srt/models/dots3_common/dots_omni_towers.py @@ -0,0 +1,240 @@ +"""In-process vision/audio towers for dots.note.omni.""" + +from __future__ import annotations + +import json +import math +from collections.abc import Iterable +from pathlib import Path + +import numpy as np +import torch +from PIL import Image + +from sglang.srt.models.dots3_common.dots_omni_audio import ( + OmniAudioConfig, + OmniAudioModel, + compute_audio_token_length, +) +from sglang.srt.models.dots3_common.dots_omni_vision import ( + DotsMoEVitConfig, + DotsMoEVitModel, +) + + +def _read_json(path: Path) -> dict: + with path.open() as file: + return json.load(file) + + +def load_omni_component_config(model_dir: Path, component: str) -> dict: + """Read a tower config from a flat dots.note.omni publish.""" + config = _read_json(model_dir / "config.json") + nested_name = f"{component}_config" + if nested_name not in config: + raise KeyError(f"Missing {nested_name!r} in {model_dir / 'config.json'}") + return config[nested_name] + + +class DotsNoteOmniVisionEncoder(DotsMoEVitModel): + """Native MoE ViT used by dots.note.omni.""" + + def __init__(self, model_dir: str): + model_dir = Path(model_dir) + config = DotsMoEVitConfig(**load_omni_component_config(model_dir, "vision")) + super().__init__(config) + self.to(torch.bfloat16) + + def load_converted_state(self, state: dict[str, torch.Tensor]): + missing, unexpected = self.load_state_dict(state, strict=False) + if missing: + raise RuntimeError(f"Dots vision tower missing weights: {missing[:8]}") + if unexpected: + raise RuntimeError( + f"Dots vision tower has unexpected weights: {unexpected[:8]}" + ) + + +class DotsNoteOmniAudioEncoder(OmniAudioModel): + """Native Dots speech encoder and adapter.""" + + def __init__(self, model_dir: str): + model_dir = Path(model_dir) + config = OmniAudioConfig(**load_omni_component_config(model_dir, "audio")) + super().__init__(config) + self.to(torch.bfloat16) + + @property + def dtype(self): + return next(self.parameters()).dtype + + def load_converted_state(self, state: dict[str, torch.Tensor]): + missing, unexpected = self.load_state_dict(state, strict=True) + if missing or unexpected: + raise RuntimeError( + "Dots audio tower weight mismatch: " + f"missing={missing[:8]}, unexpected={unexpected[:8]}" + ) + + +class DotsNoteOmniImagePreprocessor: + """CPU image preprocessing matching the converted native ViT.""" + + def __init__(self, model_dir: str): + model_dir = Path(model_dir) + config = _read_json(model_dir / "preprocessor_config.json") + config = config["vision_config"] + self.min_pixels = config["min_pixels"] + self.max_pixels = config["max_pixels"] + self.patch_size = config["patch_size"] + self.temporal_patch_size = config["temporal_patch_size"] + self.merge_size = config["merge_size"] + self.pre_pixel_shuffle = config.get("pre_pixel_shuffle", True) + self.image_mean = np.asarray(config["image_mean"], dtype=np.float32) + self.image_std = np.asarray(config["image_std"], dtype=np.float32) + image_detail_path = model_dir / "image_detail.json" + self.image_detail_config = ( + _read_json(image_detail_path).get("image_details", {}) + if image_detail_path.is_file() + else {} + ) + + @staticmethod + def _round_by_factor(value: int, factor: int) -> int: + return round(value / factor) * factor + + @staticmethod + def _ceil_by_factor(value: float, factor: int) -> int: + return math.ceil(value / factor) * factor + + @staticmethod + def _floor_by_factor(value: float, factor: int) -> int: + return math.floor(value / factor) * factor + + def _resized_size( + self, + width: int, + height: int, + min_pixels: int, + max_pixels: int, + target_height=None, + target_width=None, + ): + height = target_height or height + width = target_width or width + factor = self.patch_size * self.merge_size + if min(height, width) < factor // 4: + raise ValueError( + f"Image height/width must be at least {factor // 4}, " + f"got {height}x{width}" + ) + if max(height, width) / min(height, width) > 200: + raise ValueError("Image aspect ratio must be smaller than 200") + resized_h = max(factor, self._round_by_factor(height, factor)) + resized_w = max(factor, self._round_by_factor(width, factor)) + if resized_h * resized_w > max_pixels: + beta = math.sqrt(height * width / max_pixels) + resized_h = max(factor, self._floor_by_factor(height / beta, factor)) + resized_w = max(factor, self._floor_by_factor(width / beta, factor)) + elif resized_h * resized_w < min_pixels: + beta = math.sqrt(min_pixels / (height * width)) + resized_h = self._ceil_by_factor(height * beta, factor) + resized_w = self._ceil_by_factor(width * beta, factor) + if resized_h * resized_w > max_pixels: + beta = math.sqrt(resized_h * resized_w / max_pixels) + resized_h = max(factor, self._floor_by_factor(resized_h / beta, factor)) + resized_w = max(factor, self._floor_by_factor(resized_w / beta, factor)) + return resized_h, resized_w + + def _process_image(self, image, detail="auto"): + if not isinstance(image, Image.Image): + raise TypeError(f"Expected a PIL image, got {type(image)}") + if image.mode == "RGBA": + background = Image.new("RGB", image.size, (255, 255, 255)) + background.paste(image, mask=image.getchannel("A")) + image = background + elif image.mode != "RGB": + image = image.convert("RGB") + + detail_config = self.image_detail_config.get(detail, {}) + resized_h, resized_w = self._resized_size( + *image.size, + min_pixels=detail_config.get("min_pixels", self.min_pixels), + max_pixels=detail_config.get("max_pixels", self.max_pixels), + target_height=detail_config.get("target_height"), + target_width=detail_config.get("target_width"), + ) + image = image.resize((resized_w, resized_h), Image.Resampling.BICUBIC) + array = np.asarray(image, dtype=np.float32) / 255.0 + array = (array - self.image_mean) / self.image_std + patches = array.transpose(2, 0, 1)[None] + if patches.shape[0] == 1: + patches = np.tile(patches, (self.temporal_patch_size, 1, 1, 1)) + channel = patches.shape[1] + grid_t = patches.shape[0] // self.temporal_patch_size + grid_h = resized_h // self.patch_size + grid_w = resized_w // self.patch_size + if self.pre_pixel_shuffle: + patches = patches.reshape( + grid_t, + self.temporal_patch_size, + channel, + grid_h // self.merge_size, + self.merge_size, + self.patch_size, + grid_w // self.merge_size, + self.merge_size, + self.patch_size, + ) + patches = patches.transpose(0, 3, 6, 4, 7, 2, 1, 5, 8) + else: + patches = patches.reshape( + grid_t, + self.temporal_patch_size, + channel, + grid_h, + self.patch_size, + grid_w, + self.patch_size, + ) + patches = patches.transpose(0, 3, 5, 2, 1, 4, 6) + pixel_values = torch.from_numpy( + patches.reshape( + grid_t * grid_h * grid_w, + channel * self.temporal_patch_size * self.patch_size * self.patch_size, + ) + ) + return { + "pixel_values": pixel_values, + "image_grid_thw": torch.tensor([[grid_t, grid_h, grid_w]]), + } + + def _get_image_token_str(self, token_count: int): + return "<|img|>" + "<|imgpad|>" * token_count + "<|endofimg|>" + + def process_images(self, images: Iterable, details=None): + images = list(images) + details = details or ["auto"] * len(images) + pixel_values = [] + grids = [] + token_strings = [] + for image, detail in zip(images, details): + processed = self._process_image(image, detail) + grid = processed["image_grid_thw"] + token_count = int(grid.prod().item()) // self.merge_size**2 + pixel_values.append(processed["pixel_values"]) + grids.append(grid) + token_strings.append(self._get_image_token_str(token_count)) + return pixel_values, grids, token_strings + + +def get_audio_token_string(num_samples: int, config: OmniAudioConfig) -> str: + count = compute_audio_token_length( + num_samples, + chunk_seconds=config.chunk_seconds, + conv_temporal_stride=config.conv_temporal_stride, + merge_factor=config.merge_factor, + ) + return ( + config.audio_comp_start + config.audio_comp_span * count + config.audio_comp_end + ) diff --git a/python/sglang/srt/models/dots3_common/dots_omni_vision.py b/python/sglang/srt/models/dots3_common/dots_omni_vision.py new file mode 100644 index 000000000..3eba0dbd4 --- /dev/null +++ b/python/sglang/srt/models/dots3_common/dots_omni_vision.py @@ -0,0 +1,769 @@ +import math +from typing import Any + +import torch +import torch.nn.functional as F +from torch import nn +from torch.nn import LayerNorm +from transformers.configuration_utils import PretrainedConfig +from transformers.modeling_utils import PreTrainedModel + +from sglang.srt.layers.activation import SiluAndMul +from sglang.srt.layers.attention.vision import VisionAttention as SGLVisionAttention +from sglang.srt.layers.conv import Conv2dLayer +from sglang.srt.layers.layernorm import RMSNorm + + +class VisionRotaryEmbedding(nn.Module): + """2D vision RoPE frequency table with optional caching.""" + + def __init__( + self, + dim: int, + theta: float = 10000.0, + cache_seq_len: int | None = None, + ) -> None: + super().__init__() + inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) + self.register_buffer("inv_freq", inv_freq, persistent=False) + self._cache_seq_len = cache_seq_len + if cache_seq_len is not None: + self.register_buffer( + "freqs_cache", self._compute_freqs(cache_seq_len), persistent=False + ) + + def _compute_freqs(self, seqlen: int) -> torch.Tensor: + seq = torch.arange( + seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype + ) + return torch.outer(seq, self.inv_freq) + + def forward(self, seqlen: int) -> torch.Tensor: + if self._cache_seq_len is None: + return self._compute_freqs(seqlen) + if seqlen > self.freqs_cache.shape[0]: + self.freqs_cache = self._compute_freqs(seqlen) + return self.freqs_cache[:seqlen] + + +class VisionAttention(SGLVisionAttention): + """Dots checkpoint compatibility wrapper around SGLang vision attention.""" + + def __init__(self, config: Any) -> None: + dim = config.embed_dim + super().__init__( + embed_dim=dim, + num_heads=config.num_attention_heads, + projection_size=dim, + use_qkv_parallel=True, + flatten_batch=True, + use_data_parallel=True, + qkv_bias=config.use_bias, + proj_bias=config.use_bias, + qk_normalization_by_head_size=config.use_qk_norm, + layer_norm_eps=config.rms_norm_eps, + ) + self.register_load_state_dict_pre_hook(VisionAttention._map_qkv_weight) + + @staticmethod + def _map_qkv_weight( + module: "VisionAttention", + state_dict: dict[str, torch.Tensor], + prefix: str, + *args, + ) -> None: + for suffix in ("weight", "bias"): + source = f"{prefix}qkv.{suffix}" + target = f"{prefix}qkv_proj.{suffix}" + if source in state_dict and target not in state_dict: + state_dict[target] = state_dict.pop(source) + + def forward( + self, + hidden_states: torch.Tensor, + cu_seqlens: torch.Tensor, + max_seqlen: int, + rotary_pos_emb: torch.Tensor, + ) -> torch.Tensor: + output = super().forward( + hidden_states, + cu_seqlens=cu_seqlens, + position_embeddings=(rotary_pos_emb.cos(), rotary_pos_emb.sin()), + max_seqlen=max_seqlen, + ) + return output.squeeze(0) + + +class DotsMoEVitConfig(PretrainedConfig): + model_type: str = "dots_moe_vit" + + def __init__( + self, + embed_dim: int = 1536, + hidden_size: int = 2048, + intermediate_size: int = 4224, + moe_intermediate_size: int = 2112, + num_hidden_layers: int = 42, + num_attention_heads: int = 24, + num_channels: int = 3, + patch_size: int = 14, + spatial_merge_size: int = 2, + temporal_patch_size: int = 1, + rms_norm_eps: float = 1e-5, + use_bias: bool = False, + use_qk_norm: bool = True, + attn_implementation="flash_attention_3", + initializer_range=0.02, + is_causal=False, + post_norm=True, + gradient_checkpointing=False, + pyramid_num_routed: list[int] | None = None, + capacity_factor: float = 2.0, + router_scoring_func: str = "sigmoid", + router_scale: float = 1.0, + adapter_in_dim: int = 1536, + adapter_out_dim: int = 2048, + adapter_merge_size: int = 2, + # Adapter used to merge each 2x2 patch group. + adapter_type: str = "pixel_shuffle_mlp", + # Whether input patches and RoPE positions are already 2x2-grouped. + pre_pixel_shuffle: bool = False, + # If True, use FP8 MoE implementation + enable_fp8_moe: bool = True, + **kwargs: Any, + ): + super().__init__(**kwargs) + self.embed_dim = embed_dim + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.moe_intermediate_size = moe_intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_channels = num_channels + self.patch_size = patch_size + self.spatial_merge_size = spatial_merge_size + self.temporal_patch_size = temporal_patch_size + self.rms_norm_eps = rms_norm_eps + self.use_bias = use_bias + self.use_qk_norm = use_qk_norm + self.attn_implementation = attn_implementation + self.initializer_range = initializer_range + self.is_causal = is_causal + self.post_norm = post_norm + self.gradient_checkpointing = gradient_checkpointing + self.pyramid_num_routed = pyramid_num_routed or [] + self.capacity_factor = capacity_factor + self.router_scoring_func = router_scoring_func + self.router_scale = router_scale + self.adapter_in_dim = adapter_in_dim + self.adapter_out_dim = adapter_out_dim + self.adapter_merge_size = adapter_merge_size + if adapter_type not in ("pixel_shuffle_mlp", "patch_merger"): + raise ValueError( + f"adapter_type must be 'pixel_shuffle_mlp' or 'patch_merger', got {adapter_type!r}" + ) + self.adapter_type = adapter_type + self.pre_pixel_shuffle = pre_pixel_shuffle + self.enable_fp8_moe = enable_fp8_moe + + +# ---- FFN modules ---- + + +class DotsSwiGLUFFN(nn.Module): + def __init__(self, in_features, hidden_features, bias=False): + super().__init__() + self.fc13 = nn.Linear(in_features, hidden_features * 2, bias=bias) + self.fc2 = nn.Linear(hidden_features, in_features, bias=bias) + self.act = SiluAndMul() + self.register_load_state_dict_pre_hook( + DotsSwiGLUFFN._load_fused_fc13_from_split + ) + + @staticmethod + def _load_fused_fc13_from_split( + module: "DotsSwiGLUFFN", + state_dict: dict[str, torch.Tensor], + prefix: str, + local_metadata: dict[str, Any], + strict: bool, + missing_keys: list[str], + unexpected_keys: list[str], + error_msgs: list[str], + ) -> None: + fc13_weight_key = prefix + "fc13.weight" + fc1_weight_key = prefix + "fc1.weight" + fc3_weight_key = prefix + "fc3.weight" + fc1_weight = state_dict.get(fc1_weight_key) + fc3_weight = state_dict.get(fc3_weight_key) + if fc1_weight is not None and fc3_weight is not None: + if fc13_weight_key not in state_dict: + state_dict[fc13_weight_key] = torch.cat((fc1_weight, fc3_weight), dim=0) + state_dict.pop(fc1_weight_key) + state_dict.pop(fc3_weight_key) + + fc13_bias_key = prefix + "fc13.bias" + fc1_bias_key = prefix + "fc1.bias" + fc3_bias_key = prefix + "fc3.bias" + fc1_bias = state_dict.get(fc1_bias_key) + fc3_bias = state_dict.get(fc3_bias_key) + if fc1_bias is not None and fc3_bias is not None: + if module.fc13.bias is not None and fc13_bias_key not in state_dict: + state_dict[fc13_bias_key] = torch.cat((fc1_bias, fc3_bias), dim=0) + state_dict.pop(fc1_bias_key) + state_dict.pop(fc3_bias_key) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.fc2(self.act(self.fc13(x))) + + +def _ceil_to_multiple(v: int, multiple: int) -> int: + return ((v + multiple - 1) // multiple) * multiple + + +def _per_block_cast_to_fp8_padded( + x: torch.Tensor, + *, + use_ue8m0: bool = False, + gran_k: int = 128, +) -> tuple[torch.Tensor, torch.Tensor]: + """`per_block_cast_to_fp8` with zero-padding for non-divisible 2D tensors. + + DeepGEMM block FP8 path expects block-aligned dimensions. When `x.shape` is + not divisible by `gran_k`, this helper pads zeros on both dims to the nearest + multiple and then calls `per_block_cast_to_fp8`. + """ + if x.dim() != 2: + raise ValueError(f"expected 2D tensor, got shape={tuple(x.shape)}") + if gran_k <= 0: + raise ValueError(f"gran_k must be positive, got {gran_k}") + + from deep_gemm import per_block_cast_to_fp8 + + m, n = int(x.shape[0]), int(x.shape[1]) + m_pad = _ceil_to_multiple(m, gran_k) + n_pad = _ceil_to_multiple(n, gran_k) + + if m_pad == m and n_pad == n: + return per_block_cast_to_fp8(x.contiguous(), use_ue8m0=use_ue8m0, gran_k=gran_k) + + x_pad = torch.zeros((m_pad, n_pad), dtype=x.dtype, device=x.device) + x_pad[:m, :n] = x + return per_block_cast_to_fp8(x_pad.contiguous(), use_ue8m0=use_ue8m0, gran_k=gran_k) + + +class MoESwiGLUFFN(nn.Module): + """MoE FFN with per-expert SwiGLU experts, sigmoid/softmax gating, top-k routing.""" + + def __init__(self, config: DotsMoEVitConfig, layer_number: int): + super().__init__() + self.config = config + self.layer_number = layer_number + self.hidden_size = config.embed_dim + self.num_routed = config.pyramid_num_routed[layer_number] + self.capacity_factor = config.capacity_factor + self.router_scoring_func = config.router_scoring_func + self.router_scale = config.router_scale + + self.register_buffer( + "router_bias", torch.zeros(self.num_routed, dtype=torch.float32) + ) + + self.experts = nn.ModuleList( + [ + DotsSwiGLUFFN( + self.hidden_size, config.moe_intermediate_size, bias=config.use_bias + ) + for _ in range(self.num_routed) + ] + ) + + self.gate_weight = nn.Parameter( + torch.empty((self.num_routed, self.hidden_size), dtype=torch.float32) + ) + nn.init.kaiming_uniform_(self.gate_weight, a=math.sqrt(5)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # Keep routing and top-k selection in FP32 to avoid BF16 ties. + epsilon = 1e-9 + x_flat = x.contiguous() + num_tokens = x_flat.shape[0] + + gate_logits = F.linear(x_flat.float(), self.gate_weight.float()) + + if self.router_scoring_func == "sigmoid": + gating_prob = torch.sigmoid(gate_logits) + else: + gating_prob = torch.softmax(gate_logits, dim=-1, dtype=torch.float32) + + aggregated_output = torch.zeros_like(x_flat) + aggregated_gate = torch.zeros(num_tokens, dtype=x.dtype, device=x.device) + + topk = min(int(self.capacity_factor), self.num_routed) + + gating_with_bias = gating_prob + self.router_bias.to(torch.float32).unsqueeze(0) + _, topk_indices = torch.topk(gating_with_bias, k=topk, dim=-1, sorted=False) + + routed_weights = gating_prob.gather(1, topk_indices) + if self.router_scoring_func == "sigmoid" and topk > 1: + routed_weights = routed_weights / ( + routed_weights.sum(dim=-1, keepdim=True) + epsilon + ) + routed_weights = (routed_weights * self.router_scale).to(x_flat.dtype) + + for expert_idx in range(self.num_routed): + selected_mask = topk_indices == expert_idx + if selected_mask.sum() == 0: + continue + n_idx, top = torch.where(selected_mask) + # Fancy indexing can yield non-contiguous rows; cuBLAS bf16 GEMM may then fail + # with ``CUBLAS_STATUS_INVALID_VALUE`` inside ``F.linear``. + x_selected = x_flat[n_idx].contiguous() + expert_output = self.experts[expert_idx](x_selected) + contrib = expert_output * routed_weights[n_idx, top].unsqueeze(-1) + aggregated_output[n_idx] = aggregated_output[n_idx] + contrib + aggregated_gate[n_idx] = aggregated_gate[n_idx] + routed_weights[n_idx, top] + + aggregated_output = aggregated_output / ( + aggregated_gate.unsqueeze(-1) + epsilon + ) + return aggregated_output + + +class MoESwiGLUFFNFP8(MoESwiGLUFFN): + """FP8 variant of :class:`MoESwiGLUFFN` using fused expert kernels.""" + + def __init__(self, config: DotsMoEVitConfig, layer_number: int): + super().__init__(config, layer_number) + # ``MoESwiGLUFFN`` already builds ``DotsSwiGLUFFN`` experts. + from sglang.srt.layers.moe.moe_runner import MoeRunnerConfig + + self._moe_runner_config = MoeRunnerConfig(inplace=False) + self.register_buffer("_fused_w13_fp8", None, persistent=False) + self.register_buffer("_fused_w13_scale", None, persistent=False) + self.register_buffer("_fused_w2_fp8", None, persistent=False) + self.register_buffer("_fused_w2_scale", None, persistent=False) + self.register_load_state_dict_post_hook( + MoESwiGLUFFNFP8._post_load_pack_fused_fp8 + ) + + @staticmethod + def _post_load_pack_fused_fp8( + module: "MoESwiGLUFFNFP8", _incompatible_keys + ) -> None: + module._pack_fused_fp8_weights() + + @torch.no_grad() + def _pack_fused_fp8_weights(self) -> None: + """Stack gate+up per expert, block-quantize with DeepGEMM (128×128), layout for ``fused_moe``.""" + e_list = list(self.experts) + if not e_list: + return + w13_chunks: list[torch.Tensor] = [] + s13_chunks: list[torch.Tensor] = [] + w2_chunks: list[torch.Tensor] = [] + s2_chunks: list[torch.Tensor] = [] + for ex in e_list: + # Block-quantize gate and up separately, then stack for fused MoE w13. + w1_weight, w3_weight = ex.fc13.weight.detach().chunk(2, dim=0) + w1_bf16 = w1_weight.to(torch.bfloat16) + w3_bf16 = w3_weight.to(torch.bfloat16) + q1, s1 = _per_block_cast_to_fp8_padded(w1_bf16, use_ue8m0=False, gran_k=128) + q3, s3 = _per_block_cast_to_fp8_padded(w3_bf16, use_ue8m0=False, gran_k=128) + w13_fp8 = torch.cat([q1, q3], dim=0).contiguous() + s13 = torch.cat([s1, s3], dim=0).contiguous() + w13_chunks.append(w13_fp8) + s13_chunks.append(s13) + + w2_bf16 = ex.fc2.weight.detach().to(torch.bfloat16) + q2, s2 = _per_block_cast_to_fp8_padded(w2_bf16, use_ue8m0=False, gran_k=128) + w2_chunks.append(q2.contiguous()) + s2_chunks.append(s2) + + self._fused_w13_fp8 = torch.stack(w13_chunks, dim=0).contiguous() + self._fused_w13_scale = torch.stack(s13_chunks, dim=0).contiguous() + self._fused_w2_fp8 = torch.stack(w2_chunks, dim=0).contiguous() + self._fused_w2_scale = torch.stack(s2_chunks, dim=0).contiguous() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe import fused_moe + from sglang.srt.layers.moe.topk import StandardTopKOutput + + # Keep routing decisions in FP32 to avoid BF16 top-k ties. + epsilon = 1e-9 + x_flat = x.contiguous() + + gate_logits = F.linear(x_flat.float(), self.gate_weight.float()) + + if self.router_scoring_func == "sigmoid": + gating_prob = torch.sigmoid(gate_logits) + else: + gating_prob = torch.softmax(gate_logits, dim=-1, dtype=torch.float32) + + topk = min(int(self.capacity_factor), self.num_routed) + gating_with_bias = gating_prob + self.router_bias.to(torch.float32).unsqueeze(0) + + _, topk_indices = torch.topk(gating_with_bias, k=topk, dim=-1, sorted=False) + + routed_weights = gating_prob.gather(1, topk_indices) + if self.router_scoring_func == "sigmoid" and topk > 1: + routed_weights = routed_weights / ( + routed_weights.sum(dim=-1, keepdim=True) + epsilon + ) + routed_weights = routed_weights * float(self.router_scale) + + topk_ids = topk_indices.to(torch.int32) + topk_output = StandardTopKOutput(routed_weights, topk_ids, gate_logits) + + if self._fused_w13_fp8 is None: + self._pack_fused_fp8_weights() + + b1 = b2 = None + if self.config.use_bias: + b1_list = [] + b2_list = [] + for ex in self.experts: + b1_list.append(ex.fc13.bias.detach().to(x.dtype)) + b2_list.append(ex.fc2.bias.detach().to(x.dtype)) + b1 = torch.stack(b1_list, dim=0).contiguous() + b2 = torch.stack(b2_list, dim=0).contiguous() + + fused_out = fused_moe( + x_flat, + self._fused_w13_fp8, + self._fused_w2_fp8, + topk_output, + moe_runner_config=self._moe_runner_config, + b1=b1, + b2=b2, + use_fp8_w8a8=True, + w1_scale=self._fused_w13_scale, + w2_scale=self._fused_w2_scale, + block_shape=[128, 128], + ) + denom = routed_weights.sum(dim=-1, keepdim=True).clamp_min(epsilon) + return (fused_out / denom).type_as(x) + + +# ---- PatchEmbed ---- + + +class DotsPatchEmbed(nn.Module): + def __init__(self, config: DotsMoEVitConfig): + super().__init__() + self.num_channels = config.num_channels + self.patch_size = config.patch_size + self.temporal_patch_size = config.temporal_patch_size + self.embed_dim = config.embed_dim + self.proj = Conv2dLayer( + config.num_channels, + config.embed_dim, + kernel_size=(config.patch_size, config.patch_size), + stride=(config.patch_size, config.patch_size), + ) + self.norm = RMSNorm(config.embed_dim, eps=config.rms_norm_eps) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = x.view( + -1, + self.num_channels, + self.temporal_patch_size, + self.patch_size, + self.patch_size, + )[:, :, 0] + x = self.proj(x).view(-1, self.embed_dim) + x = self.norm(x) + return x + + +# ---- Block ---- + + +class MoEVisionBlock(nn.Module): + def __init__(self, config: DotsMoEVitConfig, layer_number: int): + super().__init__() + self.attn = VisionAttention(config) + self.norm_1 = RMSNorm(config.embed_dim, eps=config.rms_norm_eps) + self.norm_2 = RMSNorm(config.embed_dim, eps=config.rms_norm_eps) + + is_moe = ( + config.pyramid_num_routed + and layer_number < len(config.pyramid_num_routed) + and config.pyramid_num_routed[layer_number] > 0 + ) + if is_moe and config.enable_fp8_moe: + self.mlp = MoESwiGLUFFNFP8(config, layer_number) + elif is_moe: + self.mlp = MoESwiGLUFFN(config, layer_number) + else: + self.mlp = DotsSwiGLUFFN( + config.embed_dim, config.intermediate_size, bias=config.use_bias + ) + + def forward( + self, + hidden_states, + cu_seqlens, + rotary_pos_emb, + max_seqlen: int, + ) -> torch.Tensor: + hidden_states = hidden_states + self.attn( + self.norm_1(hidden_states), cu_seqlens, max_seqlen, rotary_pos_emb + ) + hidden_states = hidden_states + self.mlp(self.norm_2(hidden_states)) + return hidden_states + + +# ---- Adapter (pixel_shuffle + MLP) ---- + + +def _pixel_shuffle(x, scale_factor=0.5): + if x.size(1) % 2 == 1: + x = torch.cat([x[:, :1], x], dim=1) + if x.size(2) % 2 == 1: + x = torch.cat([x[:, :, :1], x], dim=2) + n, h, w, c = x.size() + x = x.reshape(n, h, int(w * scale_factor), int(c / scale_factor)) + x = x.permute(0, 2, 1, 3).contiguous() + x = x.reshape( + n, + int(w * scale_factor), + int(h * scale_factor), + int(c / (scale_factor * scale_factor)), + ) + x = x.permute(0, 2, 1, 3).contiguous() + return x + + +class PixelShuffleAdapter(nn.Module): + """Legacy adapter: NHWC pixel-shuffle spatial merge + LayerNorm + 2-layer MLP. + + Mirrors ``cybertron`` ``FCAdapter(pool_kind='pixel_shuffle', proj_kind='mlp2x_ln_gelu')``. + State-dict keys: ``proj.0`` (LayerNorm of in_dim*merge**2), ``proj.1`` / ``proj.3`` (Linear). + """ + + def __init__(self, config: DotsMoEVitConfig): + super().__init__() + in_dim = config.adapter_in_dim + out_dim = config.adapter_out_dim + merge_size = config.adapter_merge_size + merged_dim = in_dim * merge_size**2 + self.proj = nn.Sequential( + LayerNorm(merged_dim), + nn.Linear(merged_dim, out_dim), + nn.GELU(), + nn.Linear(out_dim, out_dim), + ) + + def forward( + self, + patch_embed: torch.Tensor, + grid_thw: torch.Tensor, + ) -> torch.Tensor: + assert patch_embed.dim() == 2 and grid_thw is not None + image_features = [] + token_index = 0 + for i in range(grid_thw.shape[0]): + grid_t, grid_h, grid_w = grid_thw[i] + images_token_length = grid_t * grid_h * grid_w + _pe = patch_embed[token_index : token_index + images_token_length] + token_index += images_token_length + if grid_t == 1: + _pe = _pe.reshape(int(grid_h), int(grid_w), -1).unsqueeze(0) + else: + _pe = _pe.reshape(int(grid_t), int(grid_h), int(grid_w), -1) + _pe = _pixel_shuffle(_pe, scale_factor=0.5) + if grid_t == 1: + _pe = _pe.squeeze(0) + else: + _pe = _pe.reshape(-1, _pe.shape[-1]) + image_features.append(_pe.reshape(-1, _pe.shape[-1])) + out = torch.cat(image_features, dim=0) + out = self.proj(out) + return out + + +class PatchMergerAdapter(nn.Module): + """Cybertron ``PatchMerger`` (``pool_kind='patch_merger', proj_kind='identity'``). + + Assumes the encoder output is already laid out in ``merge_size``x``merge_size`` groups + (qwen ``pre_pixel_shuffle`` preprocessor + RoPE grouped accordingly), so merging is a + simple ``view(-1, merge**2 * in_dim)`` of consecutive tokens. State-dict layout matches + cybertron's ``PatchMerger`` (``ln_q`` over the per-token dim, ``mlp.0`` / ``mlp.2`` Linear). + """ + + def __init__(self, config: DotsMoEVitConfig): + super().__init__() + in_dim = config.adapter_in_dim + out_dim = config.adapter_out_dim + merge_size = config.adapter_merge_size + merged_dim = in_dim * merge_size**2 + self.merge_size = merge_size + self.merged_dim = merged_dim + self.ln_q = LayerNorm(in_dim, eps=1e-6) + self.mlp = nn.Sequential( + nn.Linear(merged_dim, merged_dim), + nn.GELU(), + nn.Linear(merged_dim, out_dim), + ) + + def forward( + self, + patch_embed: torch.Tensor, + grid_thw: torch.Tensor, + ) -> torch.Tensor: + assert patch_embed.dim() == 2 and grid_thw is not None + x = self.ln_q(patch_embed) + x = x.reshape(-1, self.merged_dim) + return self.mlp(x) + + +_ADAPTER_CLASSES = { + "pixel_shuffle_mlp": PixelShuffleAdapter, + "patch_merger": PatchMergerAdapter, +} + + +# ---- Full Model ---- + + +class DotsMoEVitModel(PreTrainedModel): + config_class = DotsMoEVitConfig + + def __init__(self, config: DotsMoEVitConfig) -> None: + super().__init__(config) + self.config = config + self.spatial_merge_size = config.spatial_merge_size + + self.patch_embed = DotsPatchEmbed(config) + + head_dim = config.embed_dim // config.num_attention_heads + self.rotary_pos_emb = VisionRotaryEmbedding(head_dim // 2, cache_seq_len=100000) + + self.blocks = nn.ModuleList( + [MoEVisionBlock(config, i) for i in range(config.num_hidden_layers)] + ) + + if config.post_norm: + self.post_trunk_norm = RMSNorm(config.embed_dim, eps=config.rms_norm_eps) + + adapter_cls = _ADAPTER_CLASSES.get(config.adapter_type) + if adapter_cls is None: + raise ValueError(f"Unknown adapter_type {config.adapter_type!r}") + self.adapter = adapter_cls(config) + + self.gradient_checkpointing = False + self._gradient_checkpointing_func = torch.utils.checkpoint.checkpoint + + @property + def dtype(self) -> torch.dtype: + mlp = self.blocks[0].mlp + if isinstance(mlp, DotsSwiGLUFFN): + return mlp.fc13.weight.dtype + expert = mlp.experts[0] + return expert.fc13.weight.dtype + + @property + def device(self) -> torch.device: + return self.patch_embed.proj.weight.device + + def get_pos_ids_by_grid(self, grid_thw): + # Mirrors ``cybertron`` ``AIMv2NativeModel.rot_pos_emb``: when ``pre_pixel_shuffle`` + # is set, RoPE positions follow the qwen ``merge_size`` grouped layout (default 2x2); + # otherwise positions are flat row-major regardless of ``spatial_merge_size``. + if self.config.pre_pixel_shuffle: + rope_merge_size = ( + self.spatial_merge_size if self.spatial_merge_size > 1 else 2 + ) + else: + rope_merge_size = 1 + pos_ids = [] + for t, h, w in grid_thw: + hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w) + hpos_ids = hpos_ids.reshape( + h // rope_merge_size, + rope_merge_size, + w // rope_merge_size, + rope_merge_size, + ) + hpos_ids = hpos_ids.permute(0, 2, 1, 3).flatten() + + wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1) + wpos_ids = wpos_ids.reshape( + h // rope_merge_size, + rope_merge_size, + w // rope_merge_size, + rope_merge_size, + ) + wpos_ids = wpos_ids.permute(0, 2, 1, 3).flatten() + pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1)) + return pos_ids + + def rot_pos_emb(self, grid_thw): + pos_ids = self.get_pos_ids_by_grid(grid_thw) + pos_ids = torch.cat(pos_ids, dim=0) + max_grid_size = grid_thw[:, 1:].max() + rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size) + rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1) + return rotary_pos_emb + + def _build_cu_seqlens_from_grid(self, grid_thw: torch.Tensor): + cu_seqlens = torch.repeat_interleave( + grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0] + ).cumsum( + dim=0, + dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32, + ) + cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0) + # Same as (cu_seqlens[1:] - cu_seqlens[:-1]).max(); computed here to avoid D2H inside attention. + max_seqlen = int((grid_thw[:, 1] * grid_thw[:, 2]).max().item()) + return cu_seqlens, max_seqlen + + def _build_single_temporal_cu_seqlens_from_grid(self, grid_thw: torch.Tensor): + seq_lens = grid_thw[:, 1] * grid_thw[:, 2] + cu_seqlens = torch.empty( + (grid_thw.shape[0] + 1,), device=grid_thw.device, dtype=torch.int32 + ) + cu_seqlens[0] = 0 + torch.cumsum(seq_lens, dim=0, dtype=torch.int32, out=cu_seqlens[1:]) + max_seqlen = int(seq_lens.max().item()) + return cu_seqlens, max_seqlen + + def forward( + self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, bf16=True + ) -> torch.Tensor: + if bf16: + hidden_states = hidden_states.bfloat16() + hidden_states = self.patch_embed(hidden_states) + + rotary_pos_emb = self.rot_pos_emb(grid_thw) + + if grid_thw[:, 0].sum().item() == grid_thw.shape[0]: + cu_seqlens, max_seqlen = self._build_single_temporal_cu_seqlens_from_grid( + grid_thw + ) + else: + cu_seqlens, max_seqlen = self._build_cu_seqlens_from_grid(grid_thw) + + for blk in self.blocks: + if self.gradient_checkpointing and self.training: + hidden_states = self._gradient_checkpointing_func( + blk.__call__, + hidden_states, + cu_seqlens, + rotary_pos_emb, + max_seqlen, + ) + else: + hidden_states = blk( + hidden_states, + cu_seqlens, + rotary_pos_emb, + max_seqlen, + ) + + if self.config.post_norm: + hidden_states = self.post_trunk_norm(hidden_states) + + hidden_states = self.adapter(hidden_states, grid_thw) + return hidden_states diff --git a/python/sglang/srt/models/dots3_common/fp8.py b/python/sglang/srt/models/dots3_common/fp8.py new file mode 100644 index 000000000..00692a896 --- /dev/null +++ b/python/sglang/srt/models/dots3_common/fp8.py @@ -0,0 +1,112 @@ +"""Dots-specific FP8 helpers for absorbed MLA batched matmuls.""" + +from typing import Tuple + +import torch +import triton +import triton.language as tl + +from sglang.kernels.ops.quantization.fp8_kernel import is_fp8_fnuz +from sglang.srt.utils import ceil_align + +_FP8_MAX = 224.0 if is_fp8_fnuz() else torch.finfo(torch.float8_e4m3fn).max + + +@triton.jit +def _per_token_group_quant_einsum_fp8( + x_ptr, + x_q_ptr, + x_s_ptr, + group_size, + num_b, + num_k, + total_rows, + x_stride_m, + x_stride_b, + x_q_stride_m, + x_q_stride_b, + x_s_stride_m, + x_s_stride_b, + x_s_stride_g, + eps, + quant_min, + quant_max, + BLOCK_M: tl.constexpr, + BLOCK_K: tl.constexpr, +): + row_ids = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M) + group_id = tl.program_id(1) + m_ids = row_ids // num_b + b_ids = row_ids - m_ids * num_b + k_offsets = tl.arange(0, BLOCK_K) + k_ids = group_id * group_size + k_offsets + mask = (row_ids[:, None] < total_rows) & ( + (k_offsets[None, :] < group_size) & (k_ids[None, :] < num_k) + ) + x_ptrs = ( + x_ptr + + m_ids[:, None] * x_stride_m + + b_ids[:, None] * x_stride_b + + k_ids[None, :] + ) + x = tl.load(x_ptrs, mask=mask, other=0.0).to(tl.float32) + absmax = tl.maximum(tl.max(tl.abs(x), axis=1), eps) + scale = absmax / quant_max + quant = tl.clamp(x / scale[:, None], quant_min, quant_max).to( + x_q_ptr.dtype.element_ty + ) + q_ptrs = ( + x_q_ptr + + m_ids[:, None] * x_q_stride_m + + b_ids[:, None] * x_q_stride_b + + k_ids[None, :] + ) + s_ptrs = ( + x_s_ptr + m_ids * x_s_stride_m + b_ids * x_s_stride_b + group_id * x_s_stride_g + ) + tl.store(q_ptrs, quant, mask=mask) + tl.store(s_ptrs, scale, mask=row_ids < total_rows) + + +def per_token_group_quant_einsum_fp8( + x: torch.Tensor, + group_size: int = 128, + eps: float = 1e-12, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Quantize ``[m, b, k]`` in the scale layout required by FP8 einsum.""" + assert x.ndim == 3 and x.stride(-1) == 1 + assert group_size == 128 + m, b, k = x.shape + num_groups = (k + group_size - 1) // group_size + aligned_m = ceil_align(m, 4) + x_q = x.new_empty((m, b, k), dtype=torch.float8_e4m3fn) + scale_storage = x.new_empty((b, num_groups, aligned_m), dtype=torch.float32) + x_s = scale_storage.permute(2, 0, 1)[:m] + if m == 0 or b == 0 or num_groups == 0: + return x_q, x_s + block_m = 16 + block_k = triton.next_power_of_2(group_size) + _per_token_group_quant_einsum_fp8[(triton.cdiv(m * b, block_m), num_groups)]( + x, + x_q, + x_s, + group_size, + b, + k, + m * b, + x.stride(0), + x.stride(1), + x_q.stride(0), + x_q.stride(1), + x_s.stride(0), + x_s.stride(1), + x_s.stride(2), + eps, + -_FP8_MAX, + _FP8_MAX, + block_m, + block_k, + num_warps=4, + num_stages=1, + ) + return x_q, x_s diff --git a/python/sglang/srt/models/dots3_common/modeling.py b/python/sglang/srt/models/dots3_common/modeling.py new file mode 100644 index 000000000..9a64c59a6 --- /dev/null +++ b/python/sglang/srt/models/dots3_common/modeling.py @@ -0,0 +1,2824 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +# Adapted from vLLM's DeepSeek V2 implementation: +# https://github.com/vllm-project/vllm/blob/fb6af8bc086328ca6659e72d11ffd4309ce4de22/vllm/model_executor/models/deepseek_v2.py +"""Inference implementation for dots.note.omni.""" + +import concurrent.futures +import logging +import math +from dataclasses import dataclass +from enum import IntEnum, auto +from typing import ( + Any, + Callable, + Iterable, + List, + Optional, + Protocol, + Tuple, + Union, + runtime_checkable, +) + +import torch +import torch.nn.functional as F +from torch import nn +from transformers import PretrainedConfig + +from sglang.kernels.ops.quantization.fp8_kernel import is_fp8_fnuz +from sglang.srt.batch_overlap.two_batch_overlap import ( + MaybeTboDeepEPDispatcher, + model_forward_maybe_tbo, +) +from sglang.srt.configs.dots3 import Dots3Config +from sglang.srt.distributed import ( + get_pp_group, + parallel_state, + tensor_model_parallel_all_reduce, +) +from sglang.srt.distributed.device_communicators.pynccl_allocator import ( + use_symmetric_memory, +) +from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder +from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation +from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo +from sglang.srt.layers import deep_gemm_wrapper +from sglang.srt.layers.activation import SiluAndMul +from sglang.srt.layers.attention.dsa.dsa_indexer import Indexer +from sglang.srt.layers.communicator import ( + LayerCommunicator, + LayerScatterModes, + enable_moe_dense_fully_dp, +) +from sglang.srt.layers.dp_attention import is_dp_attention_enabled +from sglang.srt.layers.layernorm import RMSNorm +from sglang.srt.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.moe import ( + get_deepep_mode, + get_moe_a2a_backend, + should_use_flashinfer_cutlass_moe_fp4_allgather, +) +from sglang.srt.layers.moe.ep_moe.layer import DeepEPMoE, get_moe_impl_class +from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE +from sglang.srt.layers.moe.topk import TopK +from sglang.srt.layers.moe.utils import is_shared_experts_fusion_disabled +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.quantization.fp8_utils import ( + block_quant_dequant, + requant_weight_ue8m0_inplace, +) +from sglang.srt.layers.radix_attention import RadixAttention +from sglang.srt.layers.rotary_embedding import get_rope_wrapper +from sglang.srt.layers.utils import PPMissingLayer, get_layer_id +from sglang.srt.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from sglang.srt.managers.mm_utils import ( + MultiModalityDataPaddingPatternTokenPairs, + general_mm_embed_routine, +) +from sglang.srt.managers.schedule_batch import ( + Modality, + MultimodalDataItem, + MultimodalInputs, +) +from sglang.srt.mem_cache.memory_pool import KVWriteLoc +from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors +from sglang.srt.model_executor.forward_context import ( + get_attn_backend, + get_token_to_kv_pool, +) +from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.models.deepseek_common.deepseek_weight_loader import ( + _load_fused_indexer_wk, +) +from sglang.srt.models.dots3_common.fp8 import per_token_group_quant_einsum_fp8 +from sglang.srt.runtime_context import ( + get_device, + get_exec, + get_parallel, +) +from sglang.srt.utils import ( + BumpAllocator, + LazyValue, + add_prefix, + bind_or_assign, + ceil_align, + ceil_div, + get_bool_env_var, + get_device_sm, + is_cuda, + is_non_idle_and_non_empty, + log_info_on_rank0, + make_layers, +) + +_is_cuda = is_cuda() +_is_fp8_fnuz = is_fp8_fnuz() +_device_sm = get_device_sm() + +# Import-time CUDA kernels would block processor imports on CPU CI. +if _is_cuda: + from sgl_kernel import merge_state_v2 + + from sglang.kernels.ops.gemm.dsv3_router_gemm import dsv3_router_gemm +else: + merge_state_v2 = None + dsv3_router_gemm = None + + +def _require_cuda() -> None: + if not _is_cuda: + raise RuntimeError("Dots3 model only supports CUDA backend.") + + +logger = logging.getLogger(__name__) + + +@runtime_checkable +class _SupportsWeightLoader(Protocol): + weight_loader: Callable[..., Any] + + +def _get_scale_block_n( + quant_config: Optional[QuantizationConfig], +) -> int: + if quant_config is not None and quant_config.weight_block_size is not None: + return quant_config.weight_block_size[0] + return 1 + + +def _get_param_weight_loader(param: nn.Parameter): + return ( + param.weight_loader + if isinstance(param, _SupportsWeightLoader) + else default_weight_loader + ) + + +class Dots3AttnForwardMethod(IntEnum): + # Use absorbed multi-latent attention + MLA = auto() + + # Use multi-head attention, but with KV cache chunked. + # This method can avoid OOM when prefix lengths are long. + MHA_CHUNKED_KV = auto() + + # Use dense MHA for short DSA prefills selected by the DSA backend. + + SWA_MHA = auto() + + +class Dots3MLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + quant_config: Optional[QuantizationConfig] = None, + reduce_results: bool = True, + prefix: str = "", + tp_rank: Optional[int] = None, + tp_size: Optional[int] = None, + ) -> None: + super().__init__() + self.tp_size = tp_size + + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + quant_config=quant_config, + prefix=add_prefix("gate_up_proj", prefix), + tp_rank=tp_rank, + tp_size=tp_size, + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + reduce_results=reduce_results, + prefix=add_prefix("down_proj", prefix), + tp_rank=tp_rank, + tp_size=tp_size, + ) + if hidden_act != "silu": + raise ValueError( + f"Unsupported activation: {hidden_act}. Only silu is supported for now." + ) + self.act_fn = SiluAndMul() + + def forward( + self, + x, + forward_batch=None, + should_allreduce_fusion: bool = False, + use_reduce_scatter: bool = False, + ): + if (self.tp_size == 1) and x.shape[0] == 0: + return x + + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj( + x, skip_all_reduce=should_allreduce_fusion or use_reduce_scatter + ) + return x + + +class Dots3MoEGate(nn.Module): + def __init__(self, config): + super().__init__() + self.weight = nn.Parameter( + torch.empty((config.n_routed_experts, config.hidden_size)) + ) + if config.topk_method == "noaux_tc": + self.e_score_correction_bias = nn.Parameter( + torch.empty((config.n_routed_experts), dtype=torch.float32) + ) + else: + self.e_score_correction_bias = None + + def forward(self, hidden_states): + # Use the fused router only for its tuned shapes. + if ( + hidden_states.shape[0] <= 16 + and hidden_states.shape[1] == 7168 + and self.weight.shape[0] == 256 + and _device_sm >= 90 + ): + # router gemm output float32 + logits = dsv3_router_gemm(hidden_states, self.weight) + else: + logits = F.linear(hidden_states, self.weight, None) + + return logits + + +class Dots3MoE(nn.Module): + def __init__( + self, + config: PretrainedConfig, + layer_id: int, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + alt_stream: Optional[torch.cuda.Stream] = None, + is_nextn: bool = False, + ): + super().__init__() + self.tp_size = get_parallel().tp_size + self.routed_scaling_factor = config.routed_scaling_factor + self.num_fused_shared_experts = ( + 0 if is_shared_experts_fusion_disabled() else config.n_shared_experts + ) + self.config = config + self.layer_id = layer_id + self.alt_stream = alt_stream + + if self.tp_size > config.n_routed_experts: + raise ValueError( + f"Tensor parallel size {self.tp_size} is greater than " + f"the number of experts {config.n_routed_experts}." + ) + + if config.hidden_act != "silu": + raise ValueError( + f"Unsupported activation: {config.hidden_act}. " + "Only silu is supported for now." + ) + + self.gate = Dots3MoEGate(config) + + self.experts = get_moe_impl_class(quant_config)( + num_experts=config.n_routed_experts + + self.num_fused_shared_experts + + get_exec().moe.ep_num_redundant_experts, + num_fused_shared_experts=self.num_fused_shared_experts, + top_k=config.num_experts_per_tok + self.num_fused_shared_experts, + hidden_size=config.hidden_size, + intermediate_size=config.moe_intermediate_size, + layer_id=self.layer_id, + quant_config=quant_config, + routed_scaling_factor=self.routed_scaling_factor, + prefix=add_prefix("experts", prefix), + ) + + self.topk = TopK( + top_k=config.num_experts_per_tok + self.num_fused_shared_experts, + layer_id=self.layer_id, + renormalize=config.norm_topk_prob, + use_grouped_topk=True, + num_expert_group=config.n_group, + num_fused_shared_experts=self.num_fused_shared_experts, + topk_group=config.topk_group, + correction_bias=self.gate.e_score_correction_bias, + routed_scaling_factor=self.routed_scaling_factor, + apply_routed_scaling_factor_on_output=self.experts.should_fuse_routed_scaling_factor_in_topk, + ) + + self.shared_experts = None + if config.n_shared_experts is not None and self.num_fused_shared_experts == 0: + intermediate_size = config.moe_intermediate_size * config.n_shared_experts + # disable tp for shared experts when enable deepep moe, or with fp4 allgather + self.shared_experts = Dots3MLP( + hidden_size=config.hidden_size, + intermediate_size=intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + reduce_results=False, + prefix=add_prefix("shared_experts", prefix), + **( + dict(tp_rank=0, tp_size=1) + if get_moe_a2a_backend().is_deepep() + or should_use_flashinfer_cutlass_moe_fp4_allgather() + else {} + ), + ) + is_packed_weight = quant_config is not None and ( + quant_config.get_name() + in { + "awq", + "awq_marlin", + "moe_wna16", + } + ) + shared_experts_is_fp8 = ( + not is_packed_weight + and self.shared_experts.gate_up_proj.weight.dtype == torch.float8_e4m3fn + ) + if shared_experts_is_fp8: + assert ( + self.shared_experts.gate_up_proj.quant_method.quant_config.weight_block_size + == self.shared_experts.down_proj.quant_method.quant_config.weight_block_size + ) + + self.top_k = config.num_experts_per_tok + + if get_moe_a2a_backend().is_deepep(): + # TODO: we will support tp < ep in the future + self.ep_size = get_parallel().moe_ep_size + self.num_experts = ( + config.n_routed_experts + get_exec().moe.ep_num_redundant_experts + ) + self.renormalize = config.norm_topk_prob + self.topk_group = config.topk_group + self.num_expert_group = config.n_group + self.correction_bias = ( + self.gate.e_score_correction_bias.data + if self.gate.e_score_correction_bias is not None + else None + ) + + self.deepep_dispatcher = MaybeTboDeepEPDispatcher( + group=parallel_state.get_tp_group().device_group, + router_topk=self.top_k, + permute_fusion=True, + num_experts=self.num_experts, + num_local_experts=config.n_routed_experts // self.tp_size, + hidden_size=config.hidden_size, + params_dtype=config.torch_dtype, + deepep_mode=get_deepep_mode(), + async_finish=True, + return_recv_hook=True, + ) + + self._enable_deepep_moe = get_moe_a2a_backend().is_deepep() + + def get_moe_weights(self): + return [ + x.data + for name, x in self.experts.named_parameters() + if name not in ["correction_bias"] + ] + + def forward( + self, + hidden_states: torch.Tensor, + forward_batch: Optional[ForwardBatch] = None, + should_allreduce_fusion: bool = False, + use_reduce_scatter: bool = False, + ) -> torch.Tensor: + if not self._enable_deepep_moe: + DUAL_STREAM_TOKEN_THRESHOLD = 1024 + if ( + self.alt_stream is not None + and self.num_fused_shared_experts == 0 + and hidden_states.shape[0] > 0 + and hidden_states.shape[0] <= DUAL_STREAM_TOKEN_THRESHOLD + ): + return self.forward_normal_dual_stream( + hidden_states, + should_allreduce_fusion, + use_reduce_scatter, + ) + else: + return self.forward_normal( + hidden_states, + should_allreduce_fusion, + use_reduce_scatter, + ) + else: + return self.forward_deepep(hidden_states, forward_batch) + + def forward_normal_dual_stream( + self, + hidden_states: torch.Tensor, + should_allreduce_fusion: bool = False, + use_reduce_scatter: bool = False, + ) -> torch.Tensor: + + current_stream = torch.cuda.current_stream() + self.alt_stream.wait_stream(current_stream) + shared_output = self._forward_shared_experts(hidden_states) + + with torch.cuda.stream(self.alt_stream): + # router_logits: (num_tokens, n_experts) + router_logits = self.gate(hidden_states) + topk_output = self.topk(hidden_states, router_logits) + final_hidden_states = self.experts(hidden_states, topk_output) + + current_stream.wait_stream(self.alt_stream) + with use_symmetric_memory(parallel_state.get_tp_group()) as sm: + final_hidden_states_out = torch.empty_like(final_hidden_states) + + torch.add(final_hidden_states, shared_output, out=final_hidden_states_out) + final_hidden_states = final_hidden_states_out + sm.tag(final_hidden_states) + if ( + self.tp_size > 1 + and not should_allreduce_fusion + and not use_reduce_scatter + and not should_use_flashinfer_cutlass_moe_fp4_allgather() + ): + final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) + return final_hidden_states + + def forward_normal( + self, + hidden_states: torch.Tensor, + should_allreduce_fusion: bool = False, + use_reduce_scatter: bool = False, + ) -> torch.Tensor: + if hidden_states.shape[0] > 0: + shared_output = self._forward_shared_experts(hidden_states) + # router_logits: (num_tokens, n_experts) + router_logits = self.gate(hidden_states) + topk_output = self.topk(hidden_states, router_logits) + else: + shared_output = None + topk_output = self.topk.empty_topk_output(hidden_states.device) + + final_hidden_states = self.experts(hidden_states, topk_output) + if shared_output is not None: + with use_symmetric_memory(parallel_state.get_tp_group()) as sm: + final_hidden_states_out = torch.empty_like(final_hidden_states) + torch.add(final_hidden_states, shared_output, out=final_hidden_states_out) + final_hidden_states = final_hidden_states_out + sm.tag(final_hidden_states) + if ( + self.tp_size > 1 + and not should_allreduce_fusion + and not use_reduce_scatter + and not should_use_flashinfer_cutlass_moe_fp4_allgather() + ): + final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) + return final_hidden_states + + def forward_deepep( + self, hidden_states: torch.Tensor, forward_batch: ForwardBatch + ) -> torch.Tensor: + shared_output = None + if hidden_states.shape[0] > 0: + # router_logits: (num_tokens, n_experts) + router_logits = self.gate(hidden_states) + shared_output = self._forward_shared_experts(hidden_states) + topk_output = self.topk( + hidden_states, + router_logits, + num_token_non_padded=forward_batch.num_token_non_padded, + expert_location_dispatch_info=ExpertLocationDispatchInfo.init_new( + layer_id=self.layer_id, + ), + ) + else: + topk_output = self.topk.empty_topk_output( + hidden_states.device, layer_id=self.layer_id + ) + + final_hidden_states = self.experts( + hidden_states=hidden_states, + topk_output=topk_output, + ) + + if shared_output is not None: + x = shared_output + if self.experts.should_fuse_routed_scaling_factor_in_topk: + x.add_(final_hidden_states) + else: + x.add_(final_hidden_states, alpha=self.routed_scaling_factor) + final_hidden_states = x + elif not self.experts.should_fuse_routed_scaling_factor_in_topk: + final_hidden_states *= self.routed_scaling_factor + + return final_hidden_states + + def _forward_shared_experts(self, hidden_states): + if self.num_fused_shared_experts == 0: + return self.shared_experts(hidden_states) + else: + return None + + def op_gate(self, state): + if is_non_idle_and_non_empty( + state.forward_batch.forward_mode, state.hidden_states_mlp_input + ): + # router_logits: (num_tokens, n_experts) + state.router_logits = self.gate(state.hidden_states_mlp_input) + else: + state.router_logits = None + + def op_shared_experts(self, state): + hidden_states_mlp_input = state.pop("hidden_states_mlp_input") + if (self.num_fused_shared_experts == 0) and is_non_idle_and_non_empty( + state.forward_batch.forward_mode, hidden_states_mlp_input + ): + state.shared_output = self.shared_experts(hidden_states_mlp_input) + else: + state.shared_output = None + + def op_select_experts(self, state): + router_logits = state.pop("router_logits") + hidden_states = state.hidden_states_mlp_input + + if router_logits is not None: + with get_global_expert_distribution_recorder().with_current_layer( + self.layer_id + ): + state.topk_weights_local, state.topk_idx_local, _ = self.topk( + hidden_states=hidden_states, + router_logits=router_logits, + num_token_non_padded=state.forward_batch.num_token_non_padded, + expert_location_dispatch_info=ExpertLocationDispatchInfo.init_new( + layer_id=self.layer_id, + ), + ) + else: + state.topk_idx_local = torch.full( + (0, self.top_k), -1, dtype=torch.int, device=hidden_states.device + ) + state.topk_weights_local = torch.empty( + (0, self.top_k), dtype=torch.float32, device=hidden_states.device + ) + + def op_dispatch_a(self, state): + if self.ep_size > 1: + self.experts.deepep_dispatcher.dispatch_a( + hidden_states=state.hidden_states_mlp_input, + topk_idx=state.pop("topk_idx_local"), + topk_weights=state.pop("topk_weights_local"), + forward_batch=state.forward_batch, + tbo_subbatch_index=state.get("tbo_subbatch_index"), + ) + + def op_dispatch_b(self, state): + if self.ep_size > 1: + with get_global_expert_distribution_recorder().with_current_layer( + self.layer_id + ): + state.dispatch_output = self.experts.deepep_dispatcher.dispatch_b( + tbo_subbatch_index=state.get("tbo_subbatch_index"), + ) + + def op_experts(self, state): + state.hidden_states_experts_output = self.experts.moe_impl( + dispatch_output=state.dispatch_output, + ) + + def op_combine_a(self, state): + if self.ep_size > 1: + self.experts.deepep_dispatcher.combine_a( + hidden_states=state.pop("hidden_states_experts_output"), + topk_idx=state.dispatch_output.topk_idx, + topk_weights=state.dispatch_output.topk_weights, + forward_batch=state.forward_batch, + tbo_subbatch_index=state.get("tbo_subbatch_index"), + ) + state.pop("dispatch_output") + + def op_combine_b(self, state): + if self.ep_size > 1: + state.hidden_states_after_combine = ( + self.experts.deepep_dispatcher.combine_b( + tbo_subbatch_index=state.get("tbo_subbatch_index"), + ) + ) + + def op_output(self, state): + final_hidden_states = state.pop("hidden_states_after_combine") + + if (shared_output := state.pop("shared_output")) is not None: + x = shared_output + x.add_(final_hidden_states, alpha=self.routed_scaling_factor) + final_hidden_states = x + else: + final_hidden_states *= self.routed_scaling_factor + + state.hidden_states_mlp_output = final_hidden_states + + +# Aligned with HF's implementation, using sliding window inclusive with the last token. +# SGLang assumes exclusive. +def get_attention_sliding_window_size(config): + return config.sliding_window_size - 1 + + +class Dots3AttentionMLA(nn.Module): + @dataclass + class Dots3AttentionMLAConfig: + attention_gate_type: str + kv_lora_rank: int + q_lora_rank: int + qk_nope_head_dim: int + qk_rope_head_dim: int + num_attention_heads: int + num_key_value_heads: int + v_head_dim: int + rope_theta: float + + @classmethod + def from_config( + cls, + config: Dots3Config, + layer_type: str, + ): + if layer_type == "sliding_attention": + return cls( + attention_gate_type=config.swa_attention_gate_type, + kv_lora_rank=config.swa_kv_lora_rank, + q_lora_rank=config.swa_q_lora_rank, + qk_nope_head_dim=config.swa_qk_nope_head_dim, + qk_rope_head_dim=config.swa_qk_rope_head_dim, + num_attention_heads=config.swa_num_attention_heads, + num_key_value_heads=config.swa_num_key_value_heads, + v_head_dim=config.swa_v_head_dim, + rope_theta=config.swa_rope_theta, + ) + else: + return cls( + attention_gate_type=config.attention_gate_type, + kv_lora_rank=config.kv_lora_rank, + q_lora_rank=config.q_lora_rank, + qk_nope_head_dim=config.qk_nope_head_dim, + qk_rope_head_dim=config.qk_rope_head_dim, + num_attention_heads=config.num_attention_heads, + num_key_value_heads=config.num_key_value_heads, + v_head_dim=config.v_head_dim, + rope_theta=config.rope_theta, + ) + + def __init__( + self, + config: Dots3Config, + quant_config: Optional[QuantizationConfig] = None, + reduce_results: bool = False, + layer_id: int = None, + prefix: str = "", + alt_stream: Optional[torch.cuda.Stream] = None, + ) -> None: + super().__init__() + + attn_tp_rank = get_parallel().attn_tp_rank + attn_tp_size = get_parallel().attn_tp_size + + self.layer_id = layer_id + self.hidden_size = config.hidden_size + + # Determine the layer type and sliding window size + layer_type = config.layer_types[layer_id] + assert layer_type in {"sliding_attention", "full_attention"} + use_sliding_window = layer_type == "sliding_attention" + self.use_swa = use_sliding_window + self.sliding_window_size = ( + get_attention_sliding_window_size(config) if use_sliding_window else -1 + ) + + # Get Attention config based on layer type. + attn_config = self.Dots3AttentionMLAConfig.from_config(config, layer_type) + self.qk_nope_head_dim = attn_config.qk_nope_head_dim + self.qk_rope_head_dim = attn_config.qk_rope_head_dim + self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim + self.v_head_dim = attn_config.v_head_dim + self.q_lora_rank = attn_config.q_lora_rank + self.kv_lora_rank = attn_config.kv_lora_rank + self.apply_mla_qkv_lora_rescale = config.apply_mla_qkv_lora_rescale + self.num_heads = attn_config.num_attention_heads + assert self.num_heads % attn_tp_size == 0 + self.num_local_heads = self.num_heads // attn_tp_size + assert ( + attn_config.num_attention_heads == attn_config.num_key_value_heads + ), "Dots3 Only supports equal number of query and key value heads." + self.attention_gate_type = attn_config.attention_gate_type + assert self.attention_gate_type in { + "headwise", + "elementwise", + }, f"Unsupported attention_gate_type: {self.attention_gate_type}. Expected 'headwise' or 'elementwise'." + self.g_proj_local_dim = self.num_local_heads * ( + 1 if self.attention_gate_type == "headwise" else self.v_head_dim + ) + self.scaling = self.qk_head_dim**-0.5 + self.rope_theta = attn_config.rope_theta + self.max_position_embeddings = config.max_position_embeddings + + # For tensor parallel attention + assert self.q_lora_rank is not None, "Dots3AttentionMLA requires q_lora_rank." + scale_block_n = _get_scale_block_n(quant_config) + self.qk_rope_head_dim_padded = ceil_align(self.qk_rope_head_dim, scale_block_n) + # NOTE(xiaozhi): For the sake of DeepGEMM kernel, we align the g_proj_local_dim to 8. + self.g_proj_local_dim_padded = ceil_align( + self.g_proj_local_dim, max(8, scale_block_n) + ) + self.use_nsa = ( + config.index_n_heads is not None + and config.index_head_dim is not None + and config.index_topk is not None + and layer_type == "full_attention" + ) + fused_qkv_out_size = ( + self.q_lora_rank + + self.kv_lora_rank + + self.qk_rope_head_dim_padded + + self.g_proj_local_dim_padded + ) + self.fused_qkv_a_g_proj_with_mqa = ReplicatedLinear( + self.hidden_size, + fused_qkv_out_size, + bias=False, + quant_config=quant_config, + prefix=add_prefix("fused_qkv_a_g_proj_with_mqa", prefix), + ) + self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps) + self.q_b_out_size = self.num_heads * self.qk_head_dim + self.q_b_proj = ColumnParallelLinear( + self.q_lora_rank, + self.q_b_out_size, + bias=False, + quant_config=quant_config, + prefix=add_prefix("q_b_proj", prefix), + tp_rank=attn_tp_rank, + tp_size=attn_tp_size, + ) + + self.kv_b_proj = ColumnParallelLinear( + self.kv_lora_rank, + self.num_heads * (self.qk_nope_head_dim + self.v_head_dim), + bias=False, + quant_config=quant_config, + prefix=add_prefix("kv_b_proj", prefix), + tp_rank=attn_tp_rank, + tp_size=attn_tp_size, + ) + # O projection. + self.o_proj = RowParallelLinear( + self.num_heads * self.v_head_dim, + self.hidden_size, + bias=False, + quant_config=quant_config, + reduce_results=reduce_results, + prefix=add_prefix("o_proj", prefix), + tp_rank=attn_tp_rank, + tp_size=attn_tp_size, + ) + self.kv_a_layernorm = RMSNorm(self.kv_lora_rank, eps=config.rms_norm_eps) + self.k_rope_only_layernorm = RMSNorm( + self.qk_rope_head_dim, eps=config.rms_norm_eps + ) + + # Transformers v5 represents the default RoPE configuration as a + # non-empty ``{"rope_type": "default", ...}`` dictionary. Pass it + # through to the shared RoPE factory, which handles both that form and + # actual scaled-RoPE configurations. + self.rope_scaling = config.rope_scaling + + self.rotary_emb = get_rope_wrapper( + head_size=self.qk_rope_head_dim, + rotary_dim=self.qk_rope_head_dim, + max_position=self.max_position_embeddings, + base=self.rope_theta, + rope_scaling=self.rope_scaling, + is_neox_style=False, + device=get_device().device, + ) + + # Optional NSA (Native Sparse Attention) indexer. + if self.use_nsa: + assert ( + self.q_lora_rank is not None + ), "Dots3 NSA requires q_lora_rank to be set in the config." + self.indexer = Indexer( + hidden_size=self.hidden_size, + index_n_heads=config.index_n_heads, + index_head_dim=config.index_head_dim, + rope_head_dim=self.qk_rope_head_dim, + index_topk=config.index_topk, + q_lora_rank=self.q_lora_rank, + max_position_embeddings=self.max_position_embeddings, + rope_theta=self.rope_theta, + scale_fmt="ue8m0", + block_size=128, + rope_scaling=self.rope_scaling, + is_neox_style=False, + prefix=add_prefix("indexer", prefix), + quant_config=quant_config, + layer_id=self.layer_id, + alt_stream=alt_stream, + ) + + self.attn_mqa = RadixAttention( + num_heads=self.num_local_heads, + head_dim=self.kv_lora_rank + self.qk_rope_head_dim, + scaling=self.scaling, + num_kv_heads=1, + layer_id=self.layer_id, + v_head_dim=self.kv_lora_rank, + quant_config=quant_config, + prefix=add_prefix("attn_mqa", prefix), + sliding_window_size=self.sliding_window_size, + ) + + self.attn_mha = RadixAttention( + num_heads=self.num_local_heads, + head_dim=self.qk_head_dim, + scaling=self.scaling, + num_kv_heads=self.num_local_heads, + layer_id=self.layer_id, + v_head_dim=self.v_head_dim, + quant_config=quant_config, + prefix=add_prefix("attn_mha", prefix), + sliding_window_size=self.sliding_window_size, + ) + + self.alt_stream = alt_stream + + self.w_kc = None + self.w_vc = None + + self.w_scale_k = None + self.w_scale_v = None + + # Attention backend used by current forward batch + self.current_attention_backend = None + + def _split_fused_qkv_a_g_proj_out( + self, x: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + q_offset = 0 + kv_offset = q_offset + self.q_lora_rank + g_offset = kv_offset + self.kv_lora_rank + self.qk_rope_head_dim_padded + q = x[..., q_offset : q_offset + self.q_lora_rank] + latent_cache = x[ + ..., kv_offset : kv_offset + self.kv_lora_rank + self.qk_rope_head_dim + ] + g = x[..., g_offset : g_offset + self.g_proj_local_dim] + return q, latent_cache, g + + def dispatch_attn_forward_method( + self, forward_batch: ForwardBatch + ) -> Dots3AttnForwardMethod: + if ( + self.use_swa + and forward_batch.forward_mode.is_context_parallel_extend() + and forward_batch.attn_cp_metadata is not None + ): + raise NotImplementedError( + "Dots3 SWA attention does not support context-parallel prefill yet. " + "Please disable --enable-prefill-cp." + ) + + backend = get_attn_backend() + from sglang.srt.layers.attention.dots_hybrid_backend import ( + DotsHybridAttnBackend, + DotsSWAMLAAttnBackend, + ) + from sglang.srt.layers.attention.hybrid_attn_backend import HybridAttnBackend + + if isinstance(backend, DotsHybridAttnBackend) and self.use_swa: + backend = backend.selected_swa_backend(forward_batch) + elif isinstance(backend, HybridAttnBackend): + backend = backend._select_backend(forward_batch.forward_mode) + if isinstance(backend, DotsSWAMLAAttnBackend): + backend = backend.selected_backend(forward_batch) + backend_name = type(backend).__name__.lower() + if "flashattention" in backend_name: + attention_backend = "fa3" + elif "flashmla" in backend_name: + attention_backend = "flashmla" + elif "triton" in backend_name: + attention_backend = "triton" + elif "sparse" in backend_name or self.use_nsa: + attention_backend = "nsa" + else: + raise NotImplementedError( + f"Unsupported Dots3 attention backend: {type(backend).__name__}" + ) + self.current_attention_backend = attention_backend + + if attention_backend == "fa3": + if self.use_swa: + # Expanded SWA-MHA is prefill-only; other modes use paged MLA. + return ( + Dots3AttnForwardMethod.MLA + if forward_batch.forward_mode.is_decode_or_idle() + or forward_batch.forward_mode.is_target_verify() + or forward_batch.forward_mode.is_draft_extend_v2() + else Dots3AttnForwardMethod.SWA_MHA + ) + if forward_batch.forward_mode.is_extend_without_speculative(): + return Dots3AttnForwardMethod.MHA_CHUNKED_KV + else: + return Dots3AttnForwardMethod.MLA + elif attention_backend in ("flashmla", "triton", "nsa"): + return Dots3AttnForwardMethod.MLA + else: + raise NotImplementedError( + f"Dots3 only supports CUDA attention backends (fa3, triton, flashmla, nsa), got: {attention_backend}" + ) + + def _apply_attention_gate(self, attn_output: torch.Tensor, g: torch.Tensor): + if attn_output.ndim != 3: + attn_output = attn_output.reshape(-1, self.num_local_heads, self.v_head_dim) + + g = torch.nn.functional.sigmoid(g) + return attn_output * ( + g.unsqueeze(-1) + if self.attention_gate_type == "headwise" + else g.reshape(-1, self.num_local_heads, self.v_head_dim) + ) + + # TODO(xiaozhi): Fuse this to post_load_weights. + def _maybe_apply_lora_rescale(self, lora_tensor: torch.Tensor, lora_dim): + if not self.apply_mla_qkv_lora_rescale: + return lora_tensor + return lora_tensor * math.sqrt(self.hidden_size / lora_dim) + + @staticmethod + def _absorbed_bmm( + lhs: torch.Tensor, + weight: torch.Tensor, + out: torch.Tensor, + weight_scale: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + if weight.dtype == torch.float8_e4m3fn: + import deep_gemm + + assert weight_scale is not None + lhs_fp8 = per_token_group_quant_einsum_fp8(lhs) + deep_gemm.fp8_einsum( + "bhr,hdr->bhd", + lhs_fp8, + (weight, weight_scale), + out, + ) + else: + torch.bmm( + lhs.transpose(0, 1), + weight, + out=out.transpose(0, 1), + ) + return out + + def _get_kv_write_loc(self, forward_batch): + swa_loc = None + if self.use_swa: + backend = get_attn_backend() + from sglang.srt.layers.attention.dots_hybrid_backend import ( + DotsHybridAttnBackend, + DotsSWAMLAAttnBackend, + ) + + if isinstance(backend, (DotsHybridAttnBackend, DotsSWAMLAAttnBackend)): + backend.maybe_rebuild_metadata_after_dp_padding(forward_batch) + if isinstance(backend, DotsHybridAttnBackend): + backend = backend.selected_swa_backend(forward_batch) + metadata = backend.forward_metadata + swa_loc = metadata.swa_out_cache_loc + if ( + swa_loc is None + or swa_loc.shape[0] != forward_batch.out_cache_loc.shape[0] + ): + if isinstance(backend, DotsSWAMLAAttnBackend): + out_cache_loc = backend.select_draft_step_out_cache_loc( + forward_batch + ) + else: + out_cache_loc = forward_batch.out_cache_loc + swa_loc = get_token_to_kv_pool().translate_loc_from_full_to_swa( + out_cache_loc + ) + return KVWriteLoc(forward_batch.out_cache_loc, swa_loc=swa_loc) + + def op_prepare(self, state): + state.attn_intermediate_state = self.forward_prepare( + positions=state.positions, + hidden_states=state.pop("hidden_states_after_comm_pre_attn"), + forward_batch=state.forward_batch, + zero_allocator=state.zero_allocator, + ) + + def op_core(self, state): + state.hidden_states_after_attn = self.forward_core( + state.pop("attn_intermediate_state") + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + zero_allocator: BumpAllocator, + ): + s = self.forward_prepare( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + zero_allocator=zero_allocator, + ) + return self.forward_core(s) + + def forward_prepare( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + zero_allocator: BumpAllocator, + ): + if hidden_states.shape[0] == 0: + assert ( + not self.o_proj.reduce_results + ), "short-circuiting allreduce will lead to hangs" + return hidden_states, None, forward_batch, None + + attn_forward_method = self.dispatch_attn_forward_method(forward_batch) + + if attn_forward_method == Dots3AttnForwardMethod.MHA_CHUNKED_KV: + inner_state = self.forward_normal_chunked_kv_prepare( + positions, hidden_states, forward_batch, zero_allocator + ) + elif attn_forward_method == Dots3AttnForwardMethod.SWA_MHA: + inner_state = self.forward_swa_mha_prepare( + positions, hidden_states, forward_batch, zero_allocator + ) + elif attn_forward_method == Dots3AttnForwardMethod.MLA: + inner_state = self.forward_absorb_prepare( + positions, hidden_states, forward_batch, zero_allocator + ) + else: + raise NotImplementedError + return None, attn_forward_method, forward_batch, inner_state + + def forward_core(self, intermediate_state): + hidden_states, attn_forward_method, _, inner_state = intermediate_state + if inner_state is None: + return hidden_states + + if attn_forward_method == Dots3AttnForwardMethod.MHA_CHUNKED_KV: + return self.forward_normal_chunked_kv_core(*inner_state) + elif attn_forward_method == Dots3AttnForwardMethod.SWA_MHA: + return self.forward_swa_mha_core(*inner_state) + elif attn_forward_method == Dots3AttnForwardMethod.MLA: + return self.forward_absorb_core(*inner_state) + else: + raise NotImplementedError + + def forward_normal_prepare( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + zero_allocator: BumpAllocator, + ): + q, latent_cache, g = self._split_fused_qkv_a_g_proj_out( + self.fused_qkv_a_g_proj_with_mqa(hidden_states)[0] + ) + q = self.q_a_layernorm(q) + q = self._maybe_apply_lora_rescale(q, self.q_lora_rank) + q = self.q_b_proj(q)[0] + q = q.view(-1, self.num_local_heads, self.qk_head_dim) + + _, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) + kv_a, _ = latent_cache.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) + latent_cache = latent_cache.unsqueeze(1) + kv_a = self.kv_a_layernorm(kv_a) + kv_a = self._maybe_apply_lora_rescale(kv_a, self.kv_lora_rank) + kv = self.kv_b_proj(kv_a)[0] + kv = kv.view(-1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim) + k_nope = kv[..., : self.qk_nope_head_dim] + v = kv[..., self.qk_nope_head_dim :] + k_pe = latent_cache[:, :, self.kv_lora_rank :] + k_pe = self.k_rope_only_layernorm(k_pe) + q_pe, k_pe = self.rotary_emb(positions, q_pe, k_pe) + q[..., self.qk_nope_head_dim :] = q_pe + k = torch.empty_like(q) + k[..., : self.qk_nope_head_dim] = k_nope + k[..., self.qk_nope_head_dim :] = k_pe + + latent_cache[:, :, : self.kv_lora_rank] = kv_a.unsqueeze(1) + latent_cache[:, :, self.kv_lora_rank :] = k_pe + + # Save latent cache + get_token_to_kv_pool().set_kv_buffer( + self.attn_mha, + self._get_kv_write_loc(forward_batch), + latent_cache, + None, + ) + + return q, k, v, g, forward_batch + + def forward_absorb_prepare( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + zero_allocator: BumpAllocator, + ): + from sglang.srt.model_executor.runner_utils.capture_mode import ( + get_is_capture_mode, + ) + + q, latent_cache, g = self._split_fused_qkv_a_g_proj_out( + self.fused_qkv_a_g_proj_with_mqa(hidden_states)[0] + ) + k_nope = latent_cache[..., : self.kv_lora_rank] + + # overlap qk norm + if self.alt_stream is not None and get_is_capture_mode(): + current_stream = torch.cuda.current_stream() + self.alt_stream.wait_stream(current_stream) + q = self.q_a_layernorm(q) + q = self._maybe_apply_lora_rescale(q, self.q_lora_rank) + with torch.cuda.stream(self.alt_stream): + k_nope = self.kv_a_layernorm(k_nope) + k_nope = self._maybe_apply_lora_rescale(k_nope, self.kv_lora_rank) + current_stream.wait_stream(self.alt_stream) + else: + q = self.q_a_layernorm(q) + q = self._maybe_apply_lora_rescale(q, self.q_lora_rank) + k_nope = self.kv_a_layernorm(k_nope) + k_nope = self._maybe_apply_lora_rescale(k_nope, self.kv_lora_rank) + + # NSA indexer owns its projections and uses the standard SGLang layer. + topk_indices = None + if self.use_nsa: + topk_indices = self.indexer( + x=hidden_states, + q_lora=q, + positions=positions, + forward_batch=forward_batch, + layer_id=self.layer_id, + ) + k_nope = k_nope.unsqueeze(1) + q = self.q_b_proj(q)[0] + q = q.view(-1, self.num_local_heads, self.qk_head_dim) + + q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) + k_pe = latent_cache[..., self.kv_lora_rank :].unsqueeze(1) + k_pe = self.k_rope_only_layernorm(k_pe) + + q_nope_out = q_nope.new_empty( + (q_nope.shape[0], self.num_local_heads, self.kv_lora_rank) + ) + self._absorbed_bmm(q_nope, self.w_kc, q_nope_out, self.w_scale_k) + + q_pe, k_pe = self.rotary_emb(positions, q_pe, k_pe) + + return ( + q_pe, + k_pe, + q_nope_out, + k_nope, + g, + forward_batch, + zero_allocator, + topk_indices, + ) + + def forward_absorb_core( + self, + q_pe, + k_pe, + q_nope_out, + k_nope, + g, + forward_batch, + zero_allocator, + topk_indices=None, + ): + attn_kwargs = {} + if topk_indices is not None: + attn_kwargs["topk_indices"] = topk_indices + + from sglang.srt.layers.attention.dots_hybrid_backend import ( + DotsHybridAttnBackend, + DotsSWAMLAAttnBackend, + ) + + backend = get_attn_backend() + if ( + self.use_swa + and ( + isinstance(backend, DotsHybridAttnBackend) + or ( + isinstance(backend, DotsSWAMLAAttnBackend) + and backend.uses_flash_attention(forward_batch) + ) + ) + and ( + forward_batch.forward_mode.is_decode_or_idle() + or forward_batch.forward_mode.is_target_verify() + or forward_batch.forward_mode.is_draft_extend_v2() + ) + ): + get_token_to_kv_pool().set_mla_kv_buffer( + self.attn_mqa, + self._get_kv_write_loc(forward_batch), + k_nope, + k_pe, + ) + q_input = torch.cat([q_nope_out, q_pe], dim=-1) + attn_output = backend.forward_swa_mla_absorbed( + q_input, self.attn_mqa, forward_batch + ) + elif self.current_attention_backend in ("fa3", "nsa") or self.use_nsa: + assert ( + not self.use_nsa + or topk_indices is not None + or forward_batch.forward_mode.is_idle() + ), "NSA attention requires topk_indices from indexer." + get_token_to_kv_pool().set_mla_kv_buffer( + self.attn_mqa, + self._get_kv_write_loc(forward_batch), + k_nope, + k_pe, + ) + attn_output = self.attn_mqa( + q_nope_out, + k_nope, + k_nope, + forward_batch, + q_rope=q_pe, + k_rope=k_pe, + save_kv_cache=False, + **attn_kwargs, + ) + elif self.current_attention_backend in {"triton", "flashmla"}: + q_input = torch.cat([q_nope_out, q_pe], dim=-1) + k_input = torch.cat([k_nope, k_pe], dim=-1) + v_input = k_nope.contiguous() + + get_token_to_kv_pool().set_mla_kv_buffer( + self.attn_mqa, + self._get_kv_write_loc(forward_batch), + k_nope, + k_pe, + ) + attn_output = self.attn_mqa( + q_input, k_input, v_input, forward_batch, save_kv_cache=False + ) + + attn_output = attn_output.view(-1, self.num_local_heads, self.kv_lora_rank) + + attn_bmm_output = attn_output.new_empty( + (attn_output.shape[0], self.num_local_heads, self.v_head_dim), + ) + self._absorbed_bmm(attn_output, self.w_vc, attn_bmm_output, self.w_scale_v) + attn_bmm_output = self._apply_attention_gate(attn_bmm_output, g) + output, _ = self.o_proj( + attn_bmm_output.reshape(-1, self.num_local_heads * self.v_head_dim) + ) + + return output + + def _chunked_prefix_attn_mha( + self, + q: torch.Tensor, + accum_output: torch.Tensor, + accum_lse: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + + assert forward_batch.num_prefix_chunks is not None + for i in range(forward_batch.num_prefix_chunks): + forward_batch.set_prefix_chunk_idx(i) + + # Fetch latent cache from memory pool with precomputed chunked kv indices + latent_cache_buf = get_token_to_kv_pool().get_key_buffer( + self.attn_mha.layer_id + ) + + latent_cache = latent_cache_buf[ + forward_batch.prefix_chunk_kv_indices[i] + ].contiguous() + + kv_a_normed, k_pe = latent_cache.split( + [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 + ) + kv_a_normed = kv_a_normed.squeeze(1).contiguous() + kv = self.kv_b_proj(kv_a_normed)[0] + + kv = kv.view( + -1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim + ) + v = kv[..., self.qk_nope_head_dim :] + k_nope = kv[..., : self.qk_nope_head_dim] + + k = torch.empty( + ( + k_nope.shape[0], + self.num_local_heads, + self.qk_nope_head_dim + self.qk_rope_head_dim, + ), + dtype=v.dtype, + device=v.device, + ) + k[..., : self.qk_nope_head_dim] = k_nope + k[..., self.qk_nope_head_dim :] = k_pe + + output, lse = self.attn_mha(q, k, v, forward_batch, save_kv_cache=False) + tmp_output = torch.empty_like(accum_output) + tmp_lse = torch.empty_like(accum_lse) + merge_state_v2(output, lse, accum_output, accum_lse, tmp_output, tmp_lse) + accum_output, accum_lse = tmp_output, tmp_lse + + return accum_output + + def forward_normal_chunked_kv_prepare( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + zero_allocator: BumpAllocator, + ): + # Materialize long prefixes in chunks to limit expanded K/V memory. + return self.forward_normal_prepare( + positions, hidden_states, forward_batch, zero_allocator + ) + + def forward_normal_chunked_kv_core(self, q, k, v, g, forward_batch): + has_extend_prefix = any(forward_batch.extend_prefix_lens_cpu) + # Only initialize the info once + if has_extend_prefix and forward_batch.num_prefix_chunks is None: + forward_batch.prepare_chunked_prefix_cache_info(q.device) + backend = get_attn_backend() + from sglang.srt.layers.attention.dots_hybrid_backend import ( + DotsHybridAttnBackend, + DotsSWAMLAAttnBackend, + ) + + if isinstance(backend, (DotsHybridAttnBackend, DotsSWAMLAAttnBackend)): + backend.init_mha_chunk_metadata(forward_batch) + + # Zero padded V rows because FA3 tiles may read past cu_seqlens_q[-1]. + real_num_tokens = forward_batch.extend_num_tokens + if real_num_tokens is not None and v.shape[0] > real_num_tokens: + v[real_num_tokens:] = 0 + + forward_batch.mha_return_lse = has_extend_prefix + # Do mha for extended part without prefix + forward_batch.set_attn_attend_prefix_cache(False) + attn_output = self.attn_mha(q, k, v, forward_batch, save_kv_cache=False) + + # Do mha attention with chunked prefix cache if there are any sequence with prefix + if has_extend_prefix: + attn_output, lse = attn_output + forward_batch.set_attn_attend_prefix_cache(True) + attn_output = self._chunked_prefix_attn_mha( + q=q, + accum_output=attn_output, + accum_lse=lse, + forward_batch=forward_batch, + ) + + attn_output = self._apply_attention_gate(attn_output, g) + attn_output = attn_output.reshape(-1, self.num_local_heads * self.v_head_dim) + output, _ = self.o_proj(attn_output) + return output + + def forward_swa_mha_prepare( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + zero_allocator: BumpAllocator, + ): + q_lora, latent_cache, g = self._split_fused_qkv_a_g_proj_out( + self.fused_qkv_a_g_proj_with_mqa(hidden_states)[0] + ) + q_lora = self.q_a_layernorm(q_lora) + q_lora = self._maybe_apply_lora_rescale(q_lora, self.q_lora_rank) + q = self.q_b_proj(q_lora)[0] + q = q.view(-1, self.num_local_heads, self.qk_head_dim) + + kv_a = latent_cache[..., : self.kv_lora_rank] + kv_a = self.kv_a_layernorm(kv_a) + kv_a = self._maybe_apply_lora_rescale(kv_a, self.kv_lora_rank) + + q_pe = q[..., -self.qk_rope_head_dim :] + latent_cache = latent_cache.unsqueeze(1) + k_pe = latent_cache[..., -self.qk_rope_head_dim :] + k_pe = self.k_rope_only_layernorm(k_pe) + q_pe, k_pe = self.rotary_emb(positions, q_pe, k_pe) + q[..., self.qk_nope_head_dim :] = q_pe + latent_cache[..., : self.kv_lora_rank] = kv_a.unsqueeze(1) + latent_cache[..., self.kv_lora_rank :] = k_pe + + # Save latent cache + get_token_to_kv_pool().set_kv_buffer( + self.attn_mha, + self._get_kv_write_loc(forward_batch), + latent_cache, + None, + ) + + # Load full latent cache + backend = get_attn_backend() + from sglang.srt.layers.attention.dots_hybrid_backend import ( + DotsHybridAttnBackend, + DotsSWAMLAAttnBackend, + ) + + assert isinstance(backend, (DotsHybridAttnBackend, DotsSWAMLAAttnBackend)) + latent_cache = backend.get_swa_mla_prefill_latent_cache( + forward_batch, self.attn_mha.layer_id + ) + kv_a, k_pe = latent_cache.split( + [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 + ) + + # `split` keeps the 576-wide latent-cache stride while this projection + # consumes only the 512-wide LoRA slice. Block-FP8 activation + # quantization requires a contiguous input. + kv = self.kv_b_proj(kv_a.contiguous())[0] + kv = kv.view(-1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim) + k_nope, v = kv.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1) + + kv_len_sum = latent_cache.shape[0] + k = torch.empty( + ( + kv_len_sum, + self.num_local_heads, + self.qk_nope_head_dim + self.qk_rope_head_dim, + ), + dtype=q.dtype, + device=q.device, + ) + k[..., : self.qk_nope_head_dim] = k_nope + k[..., self.qk_nope_head_dim :] = k_pe + + return q, k, v, g, forward_batch + + def forward_swa_mha_core(self, q, k, v, g, forward_batch): + backend = get_attn_backend() + from sglang.srt.layers.attention.dots_hybrid_backend import ( + DotsHybridAttnBackend, + DotsSWAMLAAttnBackend, + ) + + assert isinstance(backend, (DotsHybridAttnBackend, DotsSWAMLAAttnBackend)) + attn_output = backend.forward_swa_mla_expanded( + q, k, v, self.attn_mha, forward_batch + ) + attn_output = self._apply_attention_gate(attn_output, g) + attn_output = attn_output.reshape(-1, self.num_local_heads * self.v_head_dim) + return self.o_proj(attn_output)[0] + + +class Dots3DecoderLayer(nn.Module): + def __init__( + self, + config: PretrainedConfig, + layer_id: int, + quant_config: Optional[QuantizationConfig] = None, + is_nextn: bool = False, + prefix: str = "", + alt_stream: Optional[torch.cuda.Stream] = None, + ) -> None: + super().__init__() + _require_cuda() + self.hidden_size = config.hidden_size + self.config = config + self.layer_id = layer_id + self.self_attn = Dots3AttentionMLA( + config=config, + quant_config=quant_config, + layer_id=layer_id, + reduce_results=False, + prefix=add_prefix("self_attn", prefix), + alt_stream=alt_stream, + ) + self.is_layer_sparse = self._is_layer_sparse(layer_id, is_nextn=is_nextn) + is_previous_layer_sparse = self._is_layer_sparse(layer_id - 1, is_nextn=False) + is_next_layer_sparse = self._is_layer_sparse(layer_id + 1, is_nextn=False) + + self.layer_scatter_modes = LayerScatterModes.init_new( + layer_id=layer_id, + num_layers=1 if is_nextn else config.num_hidden_layers, + is_layer_sparse=self.is_layer_sparse, + is_previous_layer_sparse=is_previous_layer_sparse, + is_next_layer_sparse=is_next_layer_sparse, + ) + + if self.is_layer_sparse: + self.mlp = Dots3MoE( + config=config, + quant_config=quant_config, + prefix=add_prefix("mlp", prefix), + layer_id=self.layer_id, + alt_stream=alt_stream, + is_nextn=is_nextn, + ) + else: + if enable_moe_dense_fully_dp(): + mlp_tp_rank, mlp_tp_size = 0, 1 + else: + mlp_tp_rank, mlp_tp_size = None, None + self.mlp = Dots3MLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + prefix=add_prefix("mlp", prefix), + tp_rank=mlp_tp_rank, + tp_size=mlp_tp_size, + ) + + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + self.layer_communicator = LayerCommunicator( + layer_scatter_modes=self.layer_scatter_modes, + input_layernorm=self.input_layernorm, + post_attention_layernorm=self.post_attention_layernorm, + allow_reduce_scatter=True, + is_last_layer=( + is_nextn or (self.layer_id == self.config.num_hidden_layers - 1) + ), + ) + + def _is_layer_sparse(self, layer_id: int, is_nextn: bool) -> bool: + # The MTP block uses a dense MLP. + if is_nextn: + return False + return ( + self.config.n_routed_experts is not None + and layer_id >= self.config.first_k_dense_replace + and layer_id % self.config.moe_layer_freq == 0 + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + residual: Optional[torch.Tensor], + zero_allocator: BumpAllocator, + ) -> torch.Tensor: + + hidden_states, residual = self.layer_communicator.prepare_attn( + hidden_states, residual, forward_batch + ) + + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + zero_allocator=zero_allocator, + ) + + hidden_states, residual = self.layer_communicator.prepare_mlp( + hidden_states, residual, forward_batch + ) + + should_allreduce_fusion = ( + self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer( + forward_batch + ) + ) + + # For DP with padding, reduce scatter can be used instead of all-reduce. + use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter( + forward_batch + ) + + hidden_states = self.mlp( + hidden_states, forward_batch, should_allreduce_fusion, use_reduce_scatter + ) + + if should_allreduce_fusion: + hidden_states._sglang_needs_allreduce_fusion = True + + if not should_allreduce_fusion: + hidden_states, residual = self.layer_communicator.postprocess_layer( + hidden_states, residual, forward_batch + ) + + return hidden_states, residual + + def op_comm_prepare_attn( + self, + state, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + residual: Optional[torch.Tensor], + zero_allocator: BumpAllocator, + tbo_subbatch_index: Optional[int] = None, + ): + state.hidden_states_after_comm_pre_attn, state.residual_after_input_ln = ( + self.layer_communicator.prepare_attn(hidden_states, residual, forward_batch) + ) + state.update( + dict( + forward_batch=forward_batch, + positions=positions, + zero_allocator=zero_allocator, + tbo_subbatch_index=tbo_subbatch_index, + ) + ) + + def op_comm_prepare_mlp(self, state): + state.hidden_states_mlp_input, state.residual_after_comm_pre_mlp = ( + self.layer_communicator.prepare_mlp( + state.pop("hidden_states_after_attn"), + state.pop("residual_after_input_ln"), + state.forward_batch, + ) + ) + + def op_mlp(self, state): + hidden_states = state.pop("hidden_states_mlp_input") + if not ( + enable_moe_dense_fully_dp() + and (not self.is_layer_sparse) + and hidden_states.shape[0] == 0 + ): + state.hidden_states_mlp_output = self.mlp( + hidden_states, state.forward_batch + ) + else: + state.hidden_states_mlp_output = hidden_states + + def op_comm_postprocess_layer(self, state): + hidden_states, residual = self.layer_communicator.postprocess_layer( + state.pop("hidden_states_mlp_output"), + state.pop("residual_after_comm_pre_mlp"), + state.forward_batch, + ) + + output = dict( + positions=state.positions, + hidden_states=hidden_states, + residual=residual, + forward_batch=state.forward_batch, + zero_allocator=state.zero_allocator, + tbo_subbatch_index=state.tbo_subbatch_index, + ) + + state.clear( + expect_keys={ + "positions", + "forward_batch", + "zero_allocator", + "tbo_subbatch_index", + } + ) + return output + + +class Dots3Model(nn.Module): + fall_back_to_pt_during_load = False + + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + _require_cuda() + self.first_k_dense_replace = config.first_k_dense_replace + self.pp_group = get_pp_group() + + if self.pp_group.is_first_rank: + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + enable_tp=not is_dp_attention_enabled(), + ) + else: + self.embed_tokens = PPMissingLayer() + + self.alt_stream = torch.cuda.Stream() if _is_cuda else None + self.layers, self.start_layer, self.end_layer = make_layers( + config.num_hidden_layers, + lambda idx, prefix: Dots3DecoderLayer( + config=config, + layer_id=idx, + quant_config=quant_config, + prefix=prefix, + alt_stream=self.alt_stream, + ), + pp_rank=self.pp_group.rank_in_group, + pp_size=self.pp_group.world_size, + prefix=add_prefix("layers", prefix), + ) + if self.pp_group.is_last_rank: + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + else: + self.norm = PPMissingLayer(return_tuple=True) + + def get_input_embeddings(self) -> torch.Tensor: + return self.embed_tokens + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + pp_proxy_tensors: Optional[PPProxyTensors] = None, + ) -> Union[torch.Tensor, PPProxyTensors]: + total_num_layers = self.end_layer - self.start_layer + device = input_embeds.device if input_embeds is not None else input_ids.device + zero_allocator = BumpAllocator( + buffer_size=total_num_layers * 2 * (2 if forward_batch.can_run_tbo else 1), + dtype=torch.float32, + device=device, + ) + + if self.pp_group.is_first_rank: + if input_embeds is None: + hidden_states = self.embed_tokens(input_ids) + else: + hidden_states = input_embeds + residual = None + else: + assert pp_proxy_tensors is not None + hidden_states = pp_proxy_tensors["hidden_states"] + residual = pp_proxy_tensors["residual"] + + normal_start_layer = self.start_layer + normal_end_layer = self.end_layer + if forward_batch.can_run_tbo: + if ( + self.first_k_dense_replace > normal_start_layer + and self.first_k_dense_replace < normal_end_layer + ): + normal_end_layer = self.first_k_dense_replace + elif self.first_k_dense_replace < normal_start_layer: + normal_end_layer = normal_start_layer = 0 + + for i in range(normal_start_layer, normal_end_layer): + with get_global_expert_distribution_recorder().with_current_layer(i): + layer = self.layers[i] + hidden_states, residual = layer( + positions, hidden_states, forward_batch, residual, zero_allocator + ) + + if normal_end_layer != self.end_layer: + hidden_states, residual = model_forward_maybe_tbo( + layers=self.layers[normal_end_layer : self.end_layer], + enable_tbo=True, + positions=positions, + forward_batch=forward_batch, + hidden_states=hidden_states, + residual=residual, + input_data_scatter_mode=self.layers[ + normal_end_layer - 1 + ].layer_scatter_modes.layer_output_mode, + zero_allocator=zero_allocator, + ) + + if not self.pp_group.is_last_rank: + return PPProxyTensors( + { + "hidden_states": hidden_states, + "residual": residual, + } + ) + else: + if not forward_batch.forward_mode.is_idle(): + if residual is None: + hidden_states = self.norm(hidden_states) + else: + hidden_states, _ = self.norm(hidden_states, residual) + return hidden_states + + +class Dots3LanguageModelForCausalLM(nn.Module): + # for quark model load + packed_modules_mapping = {} + fused_shared_experts_architecture = "Dots3NoteForCausalLM" + + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + + # for quark model load + # Always fuse q_a_proj/kv_a_proj_with_mqa/g_proj when loading Dots3. + self.fuse_qkv_a_g_proj = True + assert ( + config.q_lora_rank is not None + ), "Dots3 requires q_lora_rank to enable fused_qkv_a_g_proj_with_mqa loading." + if self.fuse_qkv_a_g_proj: + self.packed_modules_mapping["fused_qkv_a_g_proj_with_mqa"] = [ + "q_a_proj", + "kv_a_proj_with_mqa", + "g_proj", + ] + + self.pp_group = get_pp_group() + self.config = config + self.tp_size = get_parallel().tp_size + self.quant_config = quant_config + self.determine_num_fused_shared_experts() + self.model = Dots3Model( + config, quant_config, prefix=add_prefix("model", prefix) + ) + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=add_prefix("lm_head", prefix), + use_attn_tp_group=get_parallel().enable_dp_lm_head, + ) + self.logits_processor = LogitsProcessor(config) + + self._routed_experts_weights_of_layer = LazyValue( + lambda: { + layer_id: layer.mlp.get_moe_weights() + for layer_id, layer in enumerate(self.model.layers) + if isinstance(layer.mlp, Dots3MoE) + } + ) + + @property + def routed_experts_weights_of_layer(self): + return self._routed_experts_weights_of_layer.value + + @classmethod + def shared_experts_fusion_disable_reason(cls, hf_config, quant_config): + """Return why shared-expert fusion is unavailable, or ``None``.""" + if ( + not _is_cuda + or torch.cuda.get_device_capability("cuda") < (8, 0) + or hf_config.architectures[0] != cls.fused_shared_experts_architecture + or hf_config.n_routed_experts != 256 + or hf_config.n_shared_experts != 1 + ): + return "Shared-expert fusion is unsupported for this Dots3 configuration." + if get_parallel().moe_ep_size > 1: + return "Dots3 shared-expert fusion is unsupported with expert parallelism." + if quant_config is not None and quant_config.get_name() == "w4afp8": + return ( + "Dots3 W4AFP8 shared and routed experts use incompatible quantization." + ) + return None + + def determine_num_fused_shared_experts(self): + self.num_fused_shared_experts = ( + 0 if is_shared_experts_fusion_disabled() else self.config.n_shared_experts + ) + + def get_input_embeddings(self) -> nn.Embedding: + return self.model.embed_tokens + + def pad_input_ids( + self, + input_ids: List[int], + mm_inputs: MultimodalInputs, + **kwargs, + ) -> List[int]: + token_pairs = [] + if mm_inputs.im_start_id is not None and mm_inputs.im_end_id is not None: + token_pairs.append((mm_inputs.im_start_id, mm_inputs.im_end_id)) + if mm_inputs.audio_start_id is not None and mm_inputs.audio_end_id is not None: + token_pairs.append((mm_inputs.audio_start_id, mm_inputs.audio_end_id)) + return MultiModalityDataPaddingPatternTokenPairs( + data_token_pairs=token_pairs or None + ).pad_input_tokens(input_ids, mm_inputs) + + def _get_precomputed_feature(self, items: List[MultimodalDataItem]) -> torch.Tensor: + embed_param = next(self.model.embed_tokens.parameters()) + features = [] + for item in items: + feature = item.precomputed_embeddings + if feature is None: + feature = item.feature + if isinstance(feature, list): + features.extend(feature) + else: + features.append(feature) + return torch.cat( + [ + feature.to( + device=embed_param.device, + dtype=embed_param.dtype, + non_blocking=True, + ) + for feature in features + ], + dim=0, + ) + + get_image_feature = _get_precomputed_feature + get_audio_feature = _get_precomputed_feature + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + pp_proxy_tensors: Optional[PPProxyTensors] = None, + ) -> torch.Tensor: + if self.pp_group.is_first_rank and input_embeds is None: + hidden_states = general_mm_embed_routine( + input_ids=input_ids, + positions=positions, + forward_batch=forward_batch, + language_model=self.model, + multimodal_model=self, + data_embedding_funcs={ + Modality.IMAGE: self.get_image_feature, + Modality.AUDIO: self.get_audio_feature, + }, + pp_proxy_tensors=pp_proxy_tensors, + ) + else: + hidden_states = self.model( + input_ids, positions, forward_batch, input_embeds, pp_proxy_tensors + ) + + if self.pp_group.is_last_rank: + return self.logits_processor( + input_ids, hidden_states, self.lm_head, forward_batch + ) + else: + return hidden_states + + @property + def start_layer(self): + return self.model.start_layer + + @property + def end_layer(self): + return self.model.end_layer + + def post_load_weights(self, is_nextn=False, weight_names=None): + + # Perform post-processing after loading weights + if is_nextn: + # K MTP heads — one self_attn per head module. + num_nextn_layers = self.config.num_nextn_predict_layers + head_attns = [ + self.model.heads[k].decoder.self_attn for k in range(num_nextn_layers) + ] + else: + if weight_names is None: + layer_ids = range(self.model.start_layer, self.model.end_layer) + else: + layer_ids = set() + for name in weight_names: + if "kv_b_proj" in name: + layer_id = int(name.split(".")[2]) + if layer_id < self.config.num_hidden_layers: + layer_ids.add(layer_id) + head_attns = None # not used for non-nextn path + + # Iterate either over base-model layer ids OR over the K MTP head self_attns. + iterator = head_attns if is_nextn else layer_ids + + for item in iterator: + self_attn = item if is_nextn else self.model.layers[item].self_attn + w = self_attn.kv_b_proj.weight + # Only two kv_b_proj formats are supported: + # 1) BF16 weights. + # 2) FP8 block-wise weights with block_size=(128, 128), consumed by DeepGEMM grouped kernels. + weight_block_size = None + block_scale = None + + if w.dtype == torch.float8_e4m3fn: + assert ( + self.quant_config is not None + and self.quant_config.weight_block_size is not None + ), "Dots3 MLA kv_b_proj only supports FP8 block quantization with weight_block_size=(128, 128)." + weight_block_size = tuple(self.quant_config.weight_block_size) + assert weight_block_size == ( + 128, + 128, + ), f"Dots3 MLA kv_b_proj only supports FP8 block_size=(128, 128), got {weight_block_size}." + block_scale = self_attn.kv_b_proj.weight_scale_inv + + if not ( + self_attn.qk_nope_head_dim % weight_block_size[0] == 0 + and self_attn.v_head_dim % weight_block_size[0] == 0 + and get_bool_env_var("SGL_USE_DEEPGEMM_BMM", "false") + ): + # NOTE(xiaozhi): At this point, we do not requant but change to use BF16. + # The issue is that when multiplying w_kc the reduction dim can not be divided + # by 128. May modify the DeepGEMM kernel to support this. + w = block_quant_dequant( + w, + block_scale, + weight_block_size, + torch.bfloat16, + ) + else: + assert ( + w.dtype == torch.bfloat16 + ), f"Dots3 MLA kv_b_proj only supports BF16 or FP8(128x128), got dtype={w.dtype}." + + w_kc, w_vc = w.unflatten( + 0, (-1, self_attn.qk_nope_head_dim + self_attn.v_head_dim) + ).split([self_attn.qk_nope_head_dim, self_attn.v_head_dim], dim=1) + if w.dtype == torch.bfloat16: + self_attn.w_kc = bind_or_assign(self_attn.w_kc, w_kc) + self_attn.w_vc = bind_or_assign(self_attn.w_vc, w_vc.transpose(1, 2)) + else: + num_tiles_k = self_attn.qk_nope_head_dim // weight_block_size[0] + num_tiles_n = self_attn.v_head_dim // weight_block_size[0] + ws_kc, ws_vc = block_scale.unflatten( + 0, (-1, (num_tiles_k + num_tiles_n)) + ).split([num_tiles_k, num_tiles_n], dim=1) + self_attn.w_scale_k = bind_or_assign( + self_attn.w_scale_k, ws_kc.transpose(1, 2).contiguous() + ) + self_attn.w_scale_v = bind_or_assign( + self_attn.w_scale_v, ws_vc.contiguous() + ) + self_attn.w_kc = bind_or_assign( + self_attn.w_kc, w_kc.transpose(1, 2).contiguous() + ) + self_attn.w_vc = bind_or_assign(self_attn.w_vc, w_vc) + + if ( + deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM + and deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0 + and self.quant_config is not None + and self.quant_config.weight_block_size is not None + ): + self._weight_requant_ue8m0(is_nextn) + + def _weight_requant_ue8m0(self, is_nextn=False): + weight_block_size = self.quant_config.weight_block_size + + moe_layers = list( + range( + self.config.first_k_dense_replace, + self.config.num_hidden_layers, + self.config.moe_layer_freq, + ) + ) + + num_nextn_layers = self.config.num_nextn_predict_layers + num_hidden_layers = ( + num_nextn_layers if is_nextn else self.config.num_hidden_layers + ) + for layer_id in range(num_hidden_layers): + if is_nextn: + layer = self.model.heads[layer_id].decoder + else: + layer = self.model.layers[layer_id] + + for module in [ + layer.self_attn.fused_qkv_a_g_proj_with_mqa, + layer.self_attn.q_b_proj, + layer.self_attn.kv_b_proj, + layer.self_attn.o_proj, + ]: + requant_weight_ue8m0_inplace( + module.weight, module.weight_scale_inv, weight_block_size + ) + + if layer_id in moe_layers or is_nextn: + if ( + isinstance(layer.mlp, Dots3MoE) + and layer.mlp.shared_experts is not None + ): + shared_experts = layer.mlp.shared_experts + for module in [ + shared_experts.gate_up_proj, + shared_experts.down_proj, + ]: + requant_weight_ue8m0_inplace( + module.weight, module.weight_scale_inv, weight_block_size + ) + + experts = layer.mlp.experts + if isinstance(experts, DeepEPMoE): + for w in [ + experts.w13_weight_fp8, + experts.w2_weight_fp8, + ]: + requant_weight_ue8m0_inplace(w[0], w[1], weight_block_size) + else: + mlp = layer.mlp + assert isinstance(mlp, Dots3MLP) + for module in [ + mlp.gate_up_proj, + mlp.down_proj, + ]: + requant_weight_ue8m0_inplace( + module.weight, module.weight_scale_inv, weight_block_size + ) + + def load_weights( + self, + weights: Iterable[Tuple[str, torch.Tensor]], + is_nextn=False, + extra_params_mapping=None, + ): + + if is_nextn: + num_nextn_layers = self.config.num_nextn_predict_layers + # compatible with old design: when the main model has only 1 layer, + # the MTP layer is at id 0; otherwise it starts at num_hidden_layers + # and spans num_nextn_layers consecutive layer ids. + base_nextn_layer_id = ( + 0 + if self.config.num_hidden_layers == 1 + else self.config.num_hidden_layers + ) + # Set of source layer prefixes for K MTP heads. + nextn_layer_ids = list( + range(base_nextn_layer_id, base_nextn_layer_id + num_nextn_layers) + ) + + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + + if extra_params_mapping is not None: + stacked_params_mapping.extend(extra_params_mapping) + + # Params for weights, fp8 weight scales, fp8 activation scales + # (param_name, weight_name, expert_id, shard_id) + expert_params_mapping = FusedMoE.make_expert_params_mapping( + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=self.config.n_routed_experts + self.num_fused_shared_experts, + ) + # Map per-expert input scales used by mixed-precision checkpoints. + if self.quant_config and self.quant_config.get_name() == "w4afp8": + expert_params_mapping += FusedMoE.make_expert_input_scale_params_mapping( + num_experts=self.config.n_routed_experts + ) + + # Always fuse q_a_proj/kv_a_proj_with_mqa/g_proj when loading Dots3. + fuse_qkv_a_g_proj = True + assert ( + self.config.q_lora_rank is not None + ), "Dots3 requires q_lora_rank to enable fused_qkv_a_g_proj_with_mqa loading." + cached_a_proj = {} if fuse_qkv_a_g_proj else None + attn_tp_rank = get_parallel().attn_tp_rank + attn_tp_size = get_parallel().attn_tp_size + + def shard_g_proj_for_attention_tp( + weight: torch.Tensor, cat_dim: int, is_scale: bool + ): + assert ( + weight.ndim > cat_dim + ), f"weight.ndim={weight.ndim}, cat_dim={cat_dim}" + dim_size = weight.shape[cat_dim] + if not is_scale: + assert dim_size % attn_tp_size == 0, ( + f"dim_size={dim_size} is not divisible by attn_tp_size=" + f"{attn_tp_size}" + ) + start_idx = (attn_tp_rank * dim_size) // attn_tp_size + end_idx = ceil_div((attn_tp_rank + 1) * dim_size, attn_tp_size) + return weight.narrow(cat_dim, start_idx, end_idx - start_idx) + + def align_tensor( + weight: torch.Tensor, + cat_dim: int, + align_size: int, + ) -> torch.Tensor: + original_size = weight.shape[cat_dim] + aligned_size = ceil_align(original_size, align_size) + if aligned_size == original_size: + return weight + pad_shape = list(weight.shape) + pad_shape[cat_dim] = aligned_size - original_size + return torch.cat([weight, weight.new_zeros(pad_shape)], dim=cat_dim) + + def _try_load_fused_qkv( + q_a_proj_name: str, + kv_a_proj_name: str, + g_proj_name: str, + param_name: str, + cat_dim: int, + ) -> bool: + if not ( + q_a_proj_name in cached_a_proj + and kv_a_proj_name in cached_a_proj + and g_proj_name in cached_a_proj + ): + return False + if param_name not in params_dict: + raise ValueError(f"{param_name} not found in params_dict.") + param = params_dict[param_name] + target_dim = param.shape[cat_dim] if param.dim() > cat_dim else None + is_scale = param_name.endswith(".weight_scale_inv") + q_a_proj_weight = cached_a_proj[q_a_proj_name] + kv_a_proj_weight = cached_a_proj[kv_a_proj_name] + g_proj_weight = cached_a_proj[g_proj_name] + g_proj_shard = shard_g_proj_for_attention_tp( + g_proj_weight, cat_dim, is_scale + ) + scale_block_n = _get_scale_block_n(self.quant_config) + kv_a_proj_weight_aligned = ( + align_tensor(kv_a_proj_weight, cat_dim, scale_block_n) + if not is_scale + else kv_a_proj_weight + ) + g_proj_shard_aligned = ( + align_tensor(g_proj_shard, cat_dim, max(8, scale_block_n)) + if not is_scale + else g_proj_shard + ) + fused_parts = [ + q_a_proj_weight, + kv_a_proj_weight_aligned, + g_proj_shard_aligned, + ] + fused_weight = torch.cat(fused_parts, dim=cat_dim) + fused_dim = fused_weight.shape[cat_dim] + if target_dim is not None and fused_dim != target_dim: + raise ValueError( + f"Cannot match fused_qkv_a_g_proj_with_mqa shape for {param_name}: " + f"target_dim={target_dim}, fused_dim={fused_dim}." + ) + weight_loader = _get_param_weight_loader(param) + futures.append(executor.submit(weight_loader, param, fused_weight)) + cached_a_proj.pop(q_a_proj_name) + cached_a_proj.pop(kv_a_proj_name) + cached_a_proj.pop(g_proj_name) + return True + + if is_nextn: + nextn_layer_prefixes = [f"model.layers.{lid}" for lid in nextn_layer_ids] + nextn_spec_weight_names = [ + "shared_head.norm", + "eh_proj", + "enorm", + "hnorm", + ] + + # Map source layer id -> head index (0..num_nextn_layers-1). + def _match_nextn_prefix(name: str): + for k, p in enumerate(nextn_layer_prefixes): + # require dot-boundary: layer 30 vs 300 + if name == p or name.startswith(p + "."): + return k, p + return None, None + + if self.num_fused_shared_experts > 0: + assert self.num_fused_shared_experts == 1 + log_info_on_rank0(logger, "Shared experts fusion optimization enabled.") + + with concurrent.futures.ThreadPoolExecutor() as executor: + futures = [] + params_dict = dict(self.named_parameters()) + pending_indexer_wk = {} + weight_names = [] + for name, loaded_weight in weights: + layer_id = get_layer_id(name) + if ( + layer_id is not None + and not is_nextn + and ( + layer_id < self.model.start_layer + or layer_id >= self.model.end_layer + ) + ): + continue + if self.num_fused_shared_experts > 0 and "mlp.shared_experts" in name: + name = name.replace( + "mlp.shared_experts", + f"mlp.experts.{self.config.n_routed_experts}", + ) + + weight_names.append(name) + + if not is_nextn: + num_nextn_layers = self.config.num_nextn_predict_layers + if num_nextn_layers > 0 and name.startswith("model.layers"): + name_list = name.split(".") + if ( + len(name_list) >= 3 + and int(name_list[2]) >= self.config.num_hidden_layers + ): + continue + # The NextN draft model owns model.mtp.* weights. + if num_nextn_layers > 0 and name.startswith("model.mtp."): + continue + else: + # Remap the MTP-specific embedding into the draft model. + if name == "model.mtp.embed_tokens.weight": + name = "model.embed_tokens.weight" + elif name == "model.embed_tokens.weight": + # The target model owns the main embedding. + continue + else: + head_idx, matched_prefix = _match_nextn_prefix(name) + if matched_prefix is None: + continue + # Remap an MTP head into the draft model's shared head. + if "shared_head.head" in name: + name = name.replace( + matched_prefix, "model.shared_head.head" + ) + param = params_dict.get(name) + if param is None: + continue + weight_loader = _get_param_weight_loader(param) + futures.append( + executor.submit(weight_loader, param, loaded_weight) + ) + continue + + is_decoder = True + # nextn-specific adapters (enorm/hnorm/eh_proj/shared_head.norm) + for weight_name in nextn_spec_weight_names: + if weight_name in name: + name = name.replace( + matched_prefix, f"model.heads.{head_idx}" + ) + is_decoder = False + break + # MTP transformer block weights — go under + # model.heads.{k}.decoder.* (with shared block for "block"/"full" + # sharing the duplicate writes converge on the same parameter). + if is_decoder: + name = name.replace( + matched_prefix, f"model.heads.{head_idx}.decoder" + ) + + if "rotary_emb.inv_freq" in name: + continue + + if ( + ".indexer.wk." in name or ".indexer.weights_proj." in name + ) and _load_fused_indexer_wk( + name, + loaded_weight, + params_dict, + pending_indexer_wk, + self.quant_config, + ): + continue + for param_name, weight_name, shard_id in stacked_params_mapping: + # Skip non-stacked layers and experts (experts handled below). + if weight_name not in name: + continue + # We have mlp.experts[0].gate_proj in the checkpoint. + # Since we handle the experts below in expert_params_mapping, + # we need to skip here BEFORE we update the name, otherwise + # name will be updated to mlp.experts[0].gate_up_proj, which + # will then be updated below in expert_params_mapping + # for mlp.experts[0].gate_gate_up_proj, which breaks load. + if ("mlp.experts." in name) and name not in params_dict: + continue + name = name.replace(weight_name, param_name) + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + param = params_dict[name] + weight_loader = param.weight_loader + futures.append( + executor.submit(weight_loader, param, loaded_weight, shard_id) + ) + break + else: + for mapping in expert_params_mapping: + param_name, weight_name, expert_id, shard_id = mapping + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + param = params_dict[name] + weight_loader = param.weight_loader + futures.append( + executor.submit( + weight_loader, + param, + loaded_weight, + name, + shard_id=shard_id, + expert_id=expert_id, + ) + ) + break + else: + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + # Skip loading embed_tokens if not first rank in pipeline parallelism + if ".embed_tokens." in name and not self.pp_group.is_first_rank: + continue + # Skip loading norm if not last rank in pipeline parallelism + if ".norm." in name and not self.pp_group.is_last_rank: + continue + if ( + fuse_qkv_a_g_proj + and ".self_attn." in name + and ( + "q_a_proj" in name + or "kv_a_proj_with_mqa" in name + or "g_proj" in name + ) + ): + cached_a_proj[name] = loaded_weight + if "q_a_proj" in name: + q_a_proj_name = name + kv_a_proj_name = name.replace( + "q_a_proj", "kv_a_proj_with_mqa" + ) + g_proj_name = name.replace("q_a_proj", "g_proj") + param_name = name.replace( + "q_a_proj", "fused_qkv_a_g_proj_with_mqa" + ) + elif "kv_a_proj_with_mqa" in name: + q_a_proj_name = name.replace( + "kv_a_proj_with_mqa", "q_a_proj" + ) + kv_a_proj_name = name + g_proj_name = name.replace( + "kv_a_proj_with_mqa", "g_proj" + ) + param_name = name.replace( + "kv_a_proj_with_mqa", "fused_qkv_a_g_proj_with_mqa" + ) + else: + q_a_proj_name = name.replace("g_proj", "q_a_proj") + kv_a_proj_name = name.replace( + "g_proj", "kv_a_proj_with_mqa" + ) + g_proj_name = name + param_name = name.replace( + "g_proj", "fused_qkv_a_g_proj_with_mqa" + ) + + cat_dim = 0 + if self.quant_config is not None and ( + self.quant_config.get_name() == "awq" + or self.quant_config.get_name() == "awq_marlin" + or self.quant_config.get_name() == "moe_wna16" + ): + cat_dim = 1 + _try_load_fused_qkv( + q_a_proj_name, + kv_a_proj_name, + g_proj_name, + param_name, + cat_dim, + ) + else: + if ( + "k_scale" in name or "v_scale" in name + ) and name not in params_dict: + # modelopt attn kv scale is named differently + for scale in ["k_scale", "v_scale"]: + if scale in name: + name = name.replace( + f"{scale[0]}_proj", "attn_mqa" + ) + break + if name not in params_dict: + # modelopt ckpt contains not needed weights for MTP module: + # model.decoder.self_attn.attn_mqa.v_scale and + # model.decoder.self_attn.attn_mqa.k_scale + logger.warning(f"{name} not found in params_dict.") + continue + param = params_dict[name] + weight_loader = _get_param_weight_loader(param) + futures.append( + executor.submit(weight_loader, param, loaded_weight) + ) + + if fuse_qkv_a_g_proj and cached_a_proj: + unresolved = sorted(cached_a_proj.keys()) + preview = ", ".join(unresolved[:6]) + extra = f" (+{len(unresolved) - 6} more)" if len(unresolved) > 6 else "" + raise ValueError( + "Unresolved fused q/kv/g projection weights while loading " + "fused_qkv_a_g_proj_with_mqa. Missing counterparts or unexpected " + f"names: {preview}{extra}" + ) + if pending_indexer_wk: + unresolved = ", ".join(sorted(pending_indexer_wk.keys())[:6]) + raise ValueError( + "Incomplete native DSA Indexer wk weights: " + unresolved + ) + # Wait for all tasks to complete and raise any exceptions. + for future in concurrent.futures.as_completed(futures): + future.result() + + self.post_load_weights(is_nextn=is_nextn, weight_names=weight_names) + + def get_embed_and_head(self): + return self.model.embed_tokens.weight, self.lm_head.weight + + def set_embed_and_head(self, embed, head): + # Share target embeddings and output head with the draft model. + del self.model.embed_tokens.weight + del self.lm_head.weight + self.model.embed_tokens.weight = embed + self.lm_head.weight = head + torch.cuda.empty_cache() + torch.cuda.synchronize() + + def get_attention_sliding_window_size(self): + return get_attention_sliding_window_size(self.config) + + @classmethod + def get_model_config_for_expert_location(cls, config): + return ModelConfigForExpertLocation( + num_layers=config.num_hidden_layers, + num_logical_experts=config.n_routed_experts, + num_groups=None, + ) + + +class DotsNoteOmniThinkerForConditionalGeneration(nn.Module): + """Dots thinker with in-process audio, vision, and language submodels.""" + + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + from pathlib import Path + + from sglang.srt.models.dots3_common.dots_omni_towers import ( + DotsNoteOmniAudioEncoder, + DotsNoteOmniVisionEncoder, + ) + + self.config = config + self.pp_group = get_pp_group() + model_dir = Path(config._name_or_path) + self.language_model = Dots3LanguageModelForCausalLM( + config, + quant_config=quant_config, + prefix=add_prefix("language_model", prefix), + ) + # Load multimodal towers only where embeddings are produced. + if self.pp_group.is_first_rank and not config.language_only: + self.audio_tower = DotsNoteOmniAudioEncoder(str(model_dir)) + self.visual = DotsNoteOmniVisionEncoder(str(model_dir)) + else: + self.audio_tower = None + self.visual = None + + @property + def model(self): + return self.language_model.model + + @property + def lm_head(self): + return self.language_model.lm_head + + @property + def logits_processor(self): + return self.language_model.logits_processor + + @property + def routed_experts_weights_of_layer(self): + return self.language_model.routed_experts_weights_of_layer + + def get_input_embeddings(self): + return self.language_model.get_input_embeddings() + + def pad_input_ids(self, input_ids, mm_inputs, **kwargs): + return self.language_model.pad_input_ids(input_ids, mm_inputs, **kwargs) + + def get_image_feature(self, items: List[MultimodalDataItem]) -> torch.Tensor: + assert self.visual is not None + pixel_values = torch.cat([item.feature for item in items], dim=0).to( + device=self.visual.device, + dtype=self.visual.dtype, + non_blocking=True, + ) + grid_thw = torch.cat([item.image_grid_thw for item in items], dim=0).to( + device=self.visual.device, + non_blocking=True, + ) + return self.visual(pixel_values, grid_thw=grid_thw) + + def get_audio_feature(self, items: List[MultimodalDataItem]) -> torch.Tensor: + assert self.audio_tower is not None + waveforms = [ + item.feature.to( + device=self.audio_tower.device, + dtype=torch.float32, + non_blocking=True, + ) + for item in items + ] + lengths = torch.tensor( + [waveform.numel() for waveform in waveforms], + dtype=torch.long, + ) + features, token_lengths = self.audio_tower(waveforms, lengths) + expected = sum( + sum(end - start + 1 for start, end in item.offsets) for item in items + ) + if features.shape[0] != expected or sum(token_lengths) != expected: + raise RuntimeError( + "Dots audio feature/token mismatch: " + f"features={features.shape[0]}, tower_lengths={token_lengths}, " + f"placeholders={expected}" + ) + return features + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + pp_proxy_tensors: Optional[PPProxyTensors] = None, + ): + if self.pp_group.is_first_rank and input_embeds is None: + hidden_states = general_mm_embed_routine( + input_ids=input_ids, + positions=positions, + forward_batch=forward_batch, + language_model=self.language_model.model, + multimodal_model=self, + data_embedding_funcs={ + Modality.IMAGE: self.get_image_feature, + Modality.AUDIO: self.get_audio_feature, + }, + pp_proxy_tensors=pp_proxy_tensors, + ) + else: + hidden_states = self.language_model.model( + input_ids, + positions, + forward_batch, + input_embeds, + pp_proxy_tensors, + ) + + if self.pp_group.is_last_rank: + return self.language_model.logits_processor( + input_ids, + hidden_states, + self.language_model.lm_head, + forward_batch, + ) + return hidden_states + + +class DotsNoteOmniForConditionalGeneration(nn.Module): + """Native dots.note.omni conditional-generation model.""" + + packed_modules_mapping = Dots3LanguageModelForCausalLM.packed_modules_mapping + fall_back_to_pt_during_load = False + + @staticmethod + def shared_experts_fusion_disable_reason(hf_config, quant_config): + return Dots3LanguageModelForCausalLM.shared_experts_fusion_disable_reason( + hf_config, quant_config + ) + + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.config = config + self.thinker = DotsNoteOmniThinkerForConditionalGeneration( + config, + quant_config=quant_config, + prefix=add_prefix("thinker", prefix), + ) + language_model = self.thinker.language_model + self.pp_group = language_model.pp_group + self.tp_size = language_model.tp_size + self.quant_config = language_model.quant_config + self.num_fused_shared_experts = language_model.num_fused_shared_experts + self.forward = self.thinker.forward + self.pad_input_ids = self.thinker.pad_input_ids + + @property + def model(self): + return self.thinker.language_model.model + + @property + def lm_head(self): + return self.thinker.language_model.lm_head + + @property + def logits_processor(self): + return self.thinker.language_model.logits_processor + + @property + def routed_experts_weights_of_layer(self): + return self.thinker.language_model.routed_experts_weights_of_layer + + @property + def start_layer(self): + return self.thinker.language_model.start_layer + + @property + def end_layer(self): + return self.thinker.language_model.end_layer + + def get_input_embeddings(self): + return self.thinker.get_input_embeddings() + + def load_weights(self, weights, *args, **kwargs): + # Partition the flat checkpoint in one pass. This avoids reading the + # large tower tensors once through the model loader and again through + # safetensors. + load_towers = self.thinker.visual is not None + vision_state = {} + audio_state = {} + + def language_weights(): + for name, weight in weights: + if name.startswith("vision_encoder."): + if load_towers: + vision_state[name.removeprefix("vision_encoder.")] = weight + continue + if name.startswith("audio_encoder."): + if load_towers: + audio_state[name.removeprefix("audio_encoder.")] = weight + continue + yield name, weight + + self.thinker.language_model.load_weights(language_weights(), *args, **kwargs) + if load_towers: + self.thinker.visual.load_converted_state(vision_state) + self.thinker.audio_tower.load_converted_state(audio_state) + + def post_load_weights(self, *args, **kwargs): + return self.thinker.language_model.post_load_weights(*args, **kwargs) + + def get_embed_and_head(self): + return self.thinker.language_model.get_embed_and_head() + + def set_embed_and_head(self, embed, head): + return self.thinker.language_model.set_embed_and_head(embed, head) + + def get_attention_sliding_window_size(self): + return self.thinker.language_model.get_attention_sliding_window_size() + + @classmethod + def get_model_config_for_expert_location(cls, config): + return Dots3LanguageModelForCausalLM.get_model_config_for_expert_location( + config + ) + + +class Dots3NoteForCausalLM(DotsNoteOmniForConditionalGeneration): + """Canonical dots.note architecture exported by the flat checkpoint.""" diff --git a/python/sglang/srt/models/dots3_common/nextn.py b/python/sglang/srt/models/dots3_common/nextn.py new file mode 100644 index 000000000..2ae34a728 --- /dev/null +++ b/python/sglang/srt/models/dots3_common/nextn.py @@ -0,0 +1,204 @@ +"""Inference-only full-sharing Dots3 MTP / NextN draft model.""" + +import logging +from collections.abc import Iterable + +import torch +from torch import nn +from transformers import PretrainedConfig + +from sglang.srt.distributed import get_pp_group +from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder +from sglang.srt.layers.dp_attention import is_dp_attention_enabled +from sglang.srt.layers.layernorm import RMSNorm +from sglang.srt.layers.linear import ReplicatedLinear +from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.models.dots3_common.modeling import ( + Dots3DecoderLayer, + Dots3LanguageModelForCausalLM, +) +from sglang.srt.runtime_context import get_parallel +from sglang.srt.utils import BumpAllocator, add_prefix + +logger = logging.getLogger(__name__) + + +class Dots3MTPHead(nn.Module): + """The single MTP layer, recursively reused by every draft step.""" + + def __init__( + self, + config: PretrainedConfig, + quant_config: QuantizationConfig | None, + prefix: str, + ) -> None: + super().__init__() + self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.eh_proj = ReplicatedLinear( + 2 * config.hidden_size, + config.hidden_size, + bias=False, + quant_config=quant_config, + prefix=add_prefix("eh_proj", prefix), + ) + self.decoder = Dots3DecoderLayer( + config, + layer_id=0, + quant_config=quant_config, + is_nextn=True, + prefix=add_prefix("decoder", prefix), + ) + self.shared_head = nn.Module() + self.shared_head.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + +class Dot3NoteModelNextN(nn.Module): + """Text-only draft model containing one full-sharing MTP layer.""" + + def __init__( + self, + config: PretrainedConfig, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + if config.num_nextn_predict_layers != 1: + raise ValueError( + "Dots3 MTP currently supports one full-sharing layer only." + ) + if list(config.layer_types) != ["sliding_attention"]: + raise ValueError("Dots3 MTP full-sharing layer must use sliding_attention.") + + self.vocab_size = config.vocab_size + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + enable_tp=not is_dp_attention_enabled(), + prefix=add_prefix("embed_tokens", prefix), + ) + # The weight loader maps the shared MTP layer to heads.0. + self.heads = nn.ModuleList( + [Dots3MTPHead(config, quant_config, add_prefix("heads.0", prefix))] + ) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor | None = None, + ) -> torch.Tensor: + device = input_embeds.device if input_embeds is not None else input_ids.device + zero_allocator = BumpAllocator( + buffer_size=2, dtype=torch.float32, device=device + ) + hidden_states = ( + self._embed_input_ids(input_ids) if input_embeds is None else input_embeds + ) + head = self.heads[0] + if hidden_states.shape[0] > 0: + hidden_states, _ = head.eh_proj( + torch.cat( + ( + head.enorm(hidden_states), + head.hnorm(forward_batch.spec_info.hidden_states), + ), + dim=-1, + ) + ) + + residual = None + with get_global_expert_distribution_recorder().disable_this_region(): + hidden_states, residual = head.decoder( + positions, hidden_states, forward_batch, residual, zero_allocator + ) + + if not forward_batch.forward_mode.is_idle(): + if residual is None: + hidden_states = head.shared_head.norm(hidden_states) + else: + hidden_states, _ = head.shared_head.norm(hidden_states, residual) + return hidden_states + + def _embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + # Multimodal sentinels use target hidden states, so clamp their unused + # draft embedding indices to the vocabulary. + return self.embed_tokens(input_ids.clamp(min=0, max=self.vocab_size - 1)) + + +class Dots3NoteForCausalLMNextN(Dots3LanguageModelForCausalLM): + """Full-sharing Dots3 MTP draft registered for NEXTN decoding.""" + + fused_shared_experts_architecture = "Dots3NoteForCausalLMNextN" + + def __init__( + self, + config: PretrainedConfig, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + nn.Module.__init__(self) + self.config = config + self.tp_size = get_parallel().tp_size + self.quant_config = quant_config + self.pp_group = get_pp_group() + self.fuse_qkv_a_g_proj = True + self.packed_modules_mapping = { + "fused_qkv_a_g_proj_with_mqa": [ + "q_a_proj", + "kv_a_proj_with_mqa", + "g_proj", + ] + } + self.determine_num_fused_shared_experts() + + self.model = Dot3NoteModelNextN( + config, quant_config, prefix=add_prefix("model", prefix) + ) + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=add_prefix("model.shared_head.head", prefix), + use_attn_tp_group=get_parallel().enable_dp_lm_head, + ) + self.logits_processor = LogitsProcessor(config) + self._mtp_loaded_embed = False + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + hidden_states = self.model(input_ids, positions, forward_batch) + return self.logits_processor( + input_ids, hidden_states, self.lm_head, forward_batch + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): + weights = list(weights) + self._mtp_loaded_embed = any( + name.startswith("model.mtp.embed_tokens.") for name, _ in weights + ) + super().load_weights(weights, is_nextn=True) + + def set_embed_and_head(self, embed, head): + # Preserve a checkpoint-provided MTP embedding; share the output head. + if not self._mtp_loaded_embed: + del self.model.embed_tokens.weight + self.model.embed_tokens.weight = embed + else: + logger.info("Keeping the checkpoint's MTP-specific input embedding.") + del self.lm_head.weight + self.lm_head.weight = head + torch.cuda.empty_cache() + torch.cuda.synchronize() diff --git a/python/sglang/srt/models/dots3_nextn.py b/python/sglang/srt/models/dots3_nextn.py new file mode 100644 index 000000000..a540206ec --- /dev/null +++ b/python/sglang/srt/models/dots3_nextn.py @@ -0,0 +1,19 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2023-2024 SGLang Team + +"""Registry entry point for the Dots3 next-N model.""" + +from sglang.srt.models.dots3_common.nextn import ( + Dot3NoteModelNextN, + Dots3MTPHead, + Dots3NoteForCausalLMNextN, +) + +EntryClass = [Dots3NoteForCausalLMNextN] + +__all__ = [ + "Dot3NoteModelNextN", + "Dots3MTPHead", + "Dots3NoteForCausalLMNextN", + "EntryClass", +] diff --git a/python/sglang/srt/multimodal/processors/dots_note_omni.py b/python/sglang/srt/multimodal/processors/dots_note_omni.py new file mode 100644 index 000000000..f54959916 --- /dev/null +++ b/python/sglang/srt/multimodal/processors/dots_note_omni.py @@ -0,0 +1,565 @@ +import asyncio +import base64 +import hashlib +import logging +import os +import re +import time +from pathlib import Path +from typing import Any, ClassVar + +import numpy as np +import torch + +from sglang.srt.managers.io_struct import GenerateReqInput +from sglang.srt.managers.schedule_batch import ( + Modality, + MultimodalDataItem, + MultimodalProcessorOutput, +) +from sglang.srt.models.dots3 import Dots3NoteForCausalLM +from sglang.srt.models.dots3_common.dots_omni_towers import ( + DotsNoteOmniImagePreprocessor, + OmniAudioConfig, + get_audio_token_string, + load_omni_component_config, +) +from sglang.srt.multimodal.processors.base_processor import ( + BaseMultimodalProcessor, + MultimodalSpecialTokens, +) +from sglang.srt.utils import VideoData, get_video_bytes + +logger = logging.getLogger(__name__) + +_VIDEO_TOKEN_RE = re.compile(r"(|)") +_EXPANDED_VIDEO_MEDIA_RE = re.compile( + r"<\|sglang_dots_video_(?P