diff --git a/python/sglang/srt/dllm/algorithm/__init__.py b/python/sglang/srt/dllm/algorithm/__init__.py new file mode 100644 index 000000000..4d4ae9a4f --- /dev/null +++ b/python/sglang/srt/dllm/algorithm/__init__.py @@ -0,0 +1,39 @@ +import importlib +import logging +import pkgutil + +from sglang.srt.dllm.config import DllmConfig + +logger = logging.getLogger(__name__) + + +def import_algorithms(): + mapping = {} + package_name = "sglang.srt.dllm.algorithm" + package = importlib.import_module(package_name) + for _, name, ispkg in pkgutil.iter_modules(package.__path__, package_name + "."): + if ispkg: + continue + try: + module = importlib.import_module(name) + except Exception as e: + logger.warning(f"Ignore import error when loading {name}: {e}") + continue + if not hasattr(module, "Algorithm"): + continue + + algo = module.Algorithm + mapping[algo.__name__] = algo + + return mapping + + +def get_algorithm(config: DllmConfig): + try: + name = config.algorithm + return algo_name_to_cls[name](config) + except: + raise RuntimeError(f"Unknown diffusion LLM algorithm: {name}") + + +algo_name_to_cls = import_algorithms() diff --git a/python/sglang/srt/dllm/algorithm/base.py b/python/sglang/srt/dllm/algorithm/base.py new file mode 100644 index 000000000..349ddf4cd --- /dev/null +++ b/python/sglang/srt/dllm/algorithm/base.py @@ -0,0 +1,18 @@ +from sglang.srt.dllm.algorithm import get_algorithm +from sglang.srt.dllm.config import DllmConfig +from sglang.srt.server_args import ServerArgs + + +class DllmAlgorithm: + + def __init__( + self, + config: DllmConfig, + ): + self.block_size = config.block_size + self.mask_id = config.mask_id + + @staticmethod + def from_server_args(server_args: ServerArgs): + config = DllmConfig.from_server_args(server_args) + return get_algorithm(config) diff --git a/python/sglang/srt/dllm/algorithm/low_confidence.py b/python/sglang/srt/dllm/algorithm/low_confidence.py new file mode 100644 index 000000000..20e73d3b5 --- /dev/null +++ b/python/sglang/srt/dllm/algorithm/low_confidence.py @@ -0,0 +1,59 @@ +from typing import Optional, Tuple, Union + +import numpy as np +import torch +import torch.nn.functional as F + +from sglang.srt.dllm.algorithm.base import DllmAlgorithm +from sglang.srt.layers.logits_processor import LogitsProcessorOutput +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.model_executor.model_runner import ModelRunner + + +class LowConfidence(DllmAlgorithm): + + def run( + self, + model_runner: ModelRunner, + forward_batch: ForwardBatch, + ) -> Tuple[ + Union[LogitsProcessorOutput, torch.Tensor], Optional[torch.Tensor], bool + ]: + mask_index = forward_batch.input_ids == self.mask_id + start = len(forward_batch.input_ids) - torch.sum(mask_index).item() + + for _ in range(self.block_size): + mask_index = forward_batch.input_ids == self.mask_id + if torch.sum(mask_index).item() == 0: + break + + logits_output, can_run_cuda_graph = model_runner.forward( + forward_batch, pp_proxy_tensors=None + ) + + x = torch.argmax(logits_output.full_logits, dim=-1) + p = torch.squeeze( + torch.gather( + F.softmax(logits_output.full_logits, dim=-1), + dim=-1, + index=torch.unsqueeze(x, -1), + ), + -1, + ) + x = torch.where(mask_index, x, forward_batch.input_ids) + confidence = torch.where(mask_index, p, -np.inf) + transfer_index = torch.zeros_like(x, dtype=torch.bool, device=x.device) + _, select_index = torch.topk(confidence, k=1) + transfer_index[select_index] = True + + forward_batch.input_ids[transfer_index] = x[transfer_index] + + logits_output, can_run_cuda_graph = model_runner.forward( + forward_batch, pp_proxy_tensors=None + ) + + next_token_ids = forward_batch.input_ids[start:] + return logits_output, next_token_ids, can_run_cuda_graph + + +Algorithm = LowConfidence diff --git a/python/sglang/srt/dllm/config.py b/python/sglang/srt/dllm/config.py new file mode 100644 index 000000000..3c00de5fb --- /dev/null +++ b/python/sglang/srt/dllm/config.py @@ -0,0 +1,40 @@ +from sglang.srt.configs.model_config import ModelConfig +from sglang.srt.server_args import ServerArgs + + +class DllmConfig: + def __init__( + self, + mask_id: int, + block_size: int, + algorithm: str, + ): + self.algorithm = algorithm + self.block_size = block_size + self.mask_id = mask_id + + @staticmethod + def from_server_args( + server_args: ServerArgs, + ): + if server_args.dllm_algorithm is None: + return None + + config = ModelConfig.from_server_args( + server_args, + model_path=server_args.model_path, + model_revision=server_args.revision, + ) + + if config.hf_config.architectures[0] == "LLaDA2MoeModelLM": + mask_id = 156895 + else: + raise RuntimeError( + f"Unknown diffusion LLM: {config.hf_config.architectures[0]}" + ) + + return DllmConfig( + algorithm=server_args.dllm_algorithm, + block_size=server_args.dllm_block_size, + mask_id=mask_id, + ) diff --git a/python/sglang/srt/layers/attention/flashinfer_backend.py b/python/sglang/srt/layers/attention/flashinfer_backend.py index e776ebac5..302d88f7c 100644 --- a/python/sglang/srt/layers/attention/flashinfer_backend.py +++ b/python/sglang/srt/layers/attention/flashinfer_backend.py @@ -126,6 +126,8 @@ class FlashInferAttnBackend(AttentionBackend): model_runner.server_args.multi_item_scoring_delimiter ) + self.is_dllm_model = model_runner.server_args.dllm_algorithm is not None + # Parse constants self.decode_use_tensor_cores = should_use_tensor_core( kv_cache_dtype=model_runner.kv_cache_dtype, @@ -766,11 +768,16 @@ class FlashInferAttnBackend(AttentionBackend): ) else: + if not self.is_dllm_model: + # TODO: design a better interface + # For other models, use causal attention for the ragged part as previously + causal = True + o1, s1 = self.prefill_wrapper_ragged.forward_return_lse( q.view(-1, layer.tp_q_head_num, layer.head_dim), k.view(-1, layer.tp_k_head_num, layer.head_dim), v.view(-1, layer.tp_v_head_num, layer.head_dim), - causal=True, + causal=causal, sm_scale=layer.scaling, logits_soft_cap=logits_soft_cap, ) diff --git a/python/sglang/srt/layers/logits_processor.py b/python/sglang/srt/layers/logits_processor.py index e2c7d2ab6..522865765 100644 --- a/python/sglang/srt/layers/logits_processor.py +++ b/python/sglang/srt/layers/logits_processor.py @@ -99,6 +99,9 @@ class LogitsProcessorOutput: ) input_token_ids_logprobs_idx: Optional[List] = None + ## Part 4: Diffusion LLM only. + full_logits: Optional[torch.Tensor] = None + @dataclasses.dataclass class LogitsMetadata: @@ -229,7 +232,11 @@ class LogitsMetadata: class LogitsProcessor(nn.Module): def __init__( - self, config, skip_all_gather: bool = False, logit_scale: Optional[float] = None + self, + config, + skip_all_gather: bool = False, + logit_scale: Optional[float] = None, + return_full_logits: bool = False, ): super().__init__() self.config = config @@ -258,6 +265,8 @@ class LogitsProcessor(nn.Module): ): self.final_logit_softcapping = None + self.return_full_logits = return_full_logits + # enable chunked logprobs processing self.enable_logprobs_chunk = envs.SGLANG_ENABLE_LOGITS_PROCESSER_CHUNK.value # chunk size for logprobs processing @@ -491,6 +500,12 @@ class LogitsProcessor(nn.Module): input_logprob_indices, device=pruned_states.device, dtype=torch.int64 ) + full_logits = ( + self._get_logits(hidden_states, lm_head, logits_metadata) + if self.return_full_logits + else None + ) + hidden_states_to_store: Optional[torch.Tensor] = None if logits_metadata.capture_hidden_mode.need_capture(): if logits_metadata.capture_hidden_mode.is_full(): @@ -529,6 +544,7 @@ class LogitsProcessor(nn.Module): # Decode mode or extend mode without return_logprob. return LogitsProcessorOutput( + full_logits=full_logits, next_token_logits=sampled_logits, hidden_states=hidden_states_to_store, ) @@ -585,6 +601,7 @@ class LogitsProcessor(nn.Module): ) return LogitsProcessorOutput( + full_logits=full_logits, next_token_logits=sampled_logits, hidden_states=hidden_states_to_store, input_token_logprobs=logprobs_result.input_token_logprobs, diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 851e07437..50bad26c2 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -2,6 +2,7 @@ from __future__ import annotations import enum +from sglang.srt.dllm.config import DllmConfig from sglang.srt.model_executor.forward_batch_info import ForwardBatch # Copyright 2023-2024 SGLang Team @@ -442,6 +443,7 @@ class Req: sampling_params: SamplingParams, return_logprob: bool = False, top_logprobs_num: int = 0, + dllm_config: Optional[DllmConfig] = None, token_ids_logprob: List[int] = None, stream: bool = False, origin_input_ids_unpadded: Optional[Tuple[int]] = None, @@ -683,6 +685,11 @@ class Req: # For Matryoshka embeddings self.dimensions = dimensions + # For diffusion LLM + self.dllm_ids = [] + self.dllm_block_offset = 0 + self.dllm_config = dllm_config + @property def seqlen(self): return len(self.origin_input_ids) + len(self.output_ids) @@ -751,8 +758,28 @@ class Req: # Whether request reached finished condition return self.finished_reason is not None + def is_dllm(self): + return self.dllm_config is not None + def init_next_round_input(self, tree_cache: Optional[BasePrefixCache] = None): - self.fill_ids = self.origin_input_ids + self.output_ids + if self.is_dllm(): + if not self.fill_ids: + self.dllm_ids = ( + self.origin_input_ids + + [ + self.dllm_config.mask_id, + ] + * self.dllm_config.block_size + ) + else: + self.dllm_block_offset += self.dllm_config.block_size + self.dllm_ids += [ + self.dllm_config.mask_id + ] * self.dllm_config.block_size + self.fill_ids = self.dllm_ids + else: + self.fill_ids = self.origin_input_ids + self.output_ids + input_len = len(self.fill_ids) # NOTE: the matched length is at most 1 less than the input length to enable logprob computation max_prefix_len = input_len - 1 @@ -1127,6 +1154,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): # hicache pointer for synchronizing data loading from CPU to GPU hicache_consumer_index: int = -1 + # Diffusion LLM + dllm_config: Optional[DllmConfig] = None + @classmethod def init_new( cls, @@ -1138,6 +1168,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): enable_overlap: bool, spec_algorithm: SpeculativeAlgorithm, chunked_req: Optional[Req] = None, + dllm_config: Optional[DllmConfig] = None, ): return_logprob = any(req.return_logprob for req in reqs) @@ -1166,6 +1197,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): return_hidden_states=any(req.return_hidden_states for req in reqs), is_prefill_only=all(req.is_prefill_only for req in reqs), chunked_req=chunked_req, + dllm_config=dllm_config, ) def batch_size(self): @@ -1174,6 +1206,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): def is_empty(self): return len(self.reqs) == 0 + def is_dllm(self): + return self.dllm_config is not None + def prepare_encoder_info_extend(self, input_ids: List[int], seq_lens: List[int]): self.encoder_lens_cpu = [] self.encoder_cached = [] @@ -1886,6 +1921,8 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): extend_input_logprob_token_ids=self.extend_input_logprob_token_ids, is_prefill_only=self.is_prefill_only, dimensions=self.dimensions, + dllm_block_offsets=[req.dllm_block_offset for req in self.reqs], + dllm_config=self.dllm_config, ) def copy(self): @@ -1999,3 +2036,7 @@ class ModelWorkerBatch: # Whether this batch is prefill-only (no token generation needed) is_prefill_only: bool = False + + # Diffusion LLM + dllm_block_offsets: Optional[List[int]] = None + dllm_config: Optional[DllmConfig] = None diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 3ca2861e5..b0c528b73 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -60,6 +60,7 @@ from sglang.srt.disaggregation.utils import ( prepare_abort, ) from sglang.srt.distributed import get_pp_group, get_world_group +from sglang.srt.dllm.config import DllmConfig from sglang.srt.environ import envs from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder from sglang.srt.layers.dp_attention import compute_dp_attention_world_info @@ -287,6 +288,9 @@ class Scheduler( # Init model config self.model_config = ModelConfig.from_server_args(server_args) + # Init diffusion LLM config + self.dllm_config = DllmConfig.from_server_args(server_args) + # Init inter-process communication self.init_sockets(server_args, port_args) @@ -449,6 +453,10 @@ class Scheduler( # Init chunked prefill self.chunked_prefill_size = server_args.chunked_prefill_size + if self.dllm_config is not None: + # We currently leverage chunked prefill to implement block diffusion + # for diffusion LLM. + self.chunked_prefill_size = self.dllm_config.block_size if self.chunked_prefill_size <= 0: # -1 means disable self.chunked_prefill_size = None self.chunked_req = None @@ -1284,6 +1292,7 @@ class Scheduler( self.metrics_collector if self.enable_metrics else None ), http_worker_ipc=recv_req.http_worker_ipc, + dllm_config=self.dllm_config, ) req.tokenizer = self.tokenizer @@ -1600,6 +1609,10 @@ class Scheduler( self.handle_embedding_request(tokenized_req) def get_next_batch_to_run(self) -> Optional[ScheduleBatch]: + if self.dllm_config is not None: + if self.chunked_req is not None and self.chunked_req.finished(): + self.chunked_req = None + # Merge the prefill batch into the running batch chunked_req_to_exclude = set() if self.chunked_req: @@ -1832,6 +1845,7 @@ class Scheduler( self.enable_overlap, self.spec_algorithm, chunked_req=self.chunked_req, + dllm_config=self.dllm_config, ) if self.enable_hierarchical_cache: # todo (zhiqiang): disable cuda graph execution if hicache loading triggered @@ -2064,7 +2078,10 @@ class Scheduler( self.process_batch_result_decode(batch, result) trace_slice_batch(RequestStage.DECODE_LOOP, batch.reqs) elif batch.forward_mode.is_extend(): - self.process_batch_result_prefill(batch, result) + if batch.is_dllm(): + self.process_batch_result_dllm(batch, result) + else: + self.process_batch_result_prefill(batch, result) elif batch.forward_mode.is_prebuilt(): self.process_batch_result_prebuilt(batch) elif batch.forward_mode.is_idle(): diff --git a/python/sglang/srt/managers/scheduler_output_processor_mixin.py b/python/sglang/srt/managers/scheduler_output_processor_mixin.py index e3d8f9668..6997c09b3 100644 --- a/python/sglang/srt/managers/scheduler_output_processor_mixin.py +++ b/python/sglang/srt/managers/scheduler_output_processor_mixin.py @@ -281,6 +281,36 @@ class SchedulerOutputProcessorMixin: return predict_tokens + def process_batch_result_dllm( + self: Scheduler, + batch: ScheduleBatch, + result: GenerationBatchResult, + ): + if result.copy_done is not None: + result.copy_done.synchronize() + + next_token_ids = result.next_token_ids.tolist() + self.num_generated_tokens += len(next_token_ids) + + self.token_to_kv_pool_allocator.free_group_begin() + + assert len(batch.reqs) == 1, "batch size is currently expected to be 1" + req = batch.reqs[0] + + for next_token_id in next_token_ids: + req.output_ids.append(next_token_id) + req.check_finished() + + if req.finished(): + release_kv_cache(req, self.tree_cache) + req.time_stats.completion_time = time.perf_counter() + break + + self.tree_cache.cache_unfinished_req(req) + + self.stream_output(batch.reqs, batch.return_logprob) + self.token_to_kv_pool_allocator.free_group_end() + def process_batch_result_decode( self: Scheduler, batch: ScheduleBatch, diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/managers/tp_worker.py index b4f18d84d..5e8bd9241 100644 --- a/python/sglang/srt/managers/tp_worker.py +++ b/python/sglang/srt/managers/tp_worker.py @@ -22,6 +22,7 @@ import torch from sglang.srt.configs.model_config import ModelConfig from sglang.srt.distributed import get_pp_group, get_world_group +from sglang.srt.dllm.algorithm.base import DllmAlgorithm from sglang.srt.managers.io_struct import ( DestroyWeightsUpdateGroupReqInput, GetWeightsByNameReqInput, @@ -234,6 +235,9 @@ class TpModelWorker(BaseTpWorker): is_draft_model=is_draft_worker, ) + if server_args.dllm_algorithm is not None: + self.dllm_algorithm = DllmAlgorithm.from_server_args(server_args) + self._model_runner = ModelRunner( model_config=self.model_config, mem_fraction_static=server_args.mem_fraction_static, @@ -340,6 +344,9 @@ class TpModelWorker(BaseTpWorker): self.model_runner.token_to_kv_pool.size, ) + def is_dllm(self): + return hasattr(self, "dllm_algorithm") + def forward_batch_generation( self, model_worker_batch: ModelWorkerBatch, @@ -368,6 +375,16 @@ class TpModelWorker(BaseTpWorker): ) if self.pp_group.is_last_rank: + if self.is_dllm(): + logits_output, next_token_ids, can_run_cuda_graph = ( + self.dllm_algorithm.run(self.model_runner, forward_batch) + ) + return GenerationBatchResult( + logits_output=logits_output, + next_token_ids=next_token_ids, + can_run_cuda_graph=can_run_cuda_graph, + ) + logits_output, can_run_cuda_graph = self.model_runner.forward( forward_batch, pp_proxy_tensors=pp_proxy_tensors, diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index a4f2e7025..e23470dd9 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -441,8 +441,17 @@ class ForwardBatch: ) return ret - # Override the positions with spec_info - if ( + # Override the positions with diffusion LLM or spec_info + if batch.dllm_config is not None: + block_size = batch.dllm_config.block_size + ret.positions = torch.tensor( + [ + [i for i in range(block_offset, block_offset + block_size)] + for block_offset in batch.dllm_block_offsets + ], + dtype=torch.int32, + ).to(device, non_blocking=True) + elif ( ret.spec_info is not None and getattr(ret.spec_info, "positions", None) is not None ): diff --git a/python/sglang/srt/models/llada2.py b/python/sglang/srt/models/llada2.py new file mode 100644 index 000000000..b89c62f44 --- /dev/null +++ b/python/sglang/srt/models/llada2.py @@ -0,0 +1,941 @@ +# coding=utf-8 +# Copyright 2023 Antgroup and The HuggingFace Inc. team. All rights reserved. +# +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX +# and OPT implementations in this library. It has been modified from its +# original forms to accommodate minor architectural differences compared +# to GPT-NeoX and OPT used by the Meta AI team that trained the model. +# +# 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. +"""SGLang LLaDA2MoeModelLM model.""" +import logging +from typing import Iterable, Optional, Tuple, Union + +import torch +import torch.nn.functional as F +from torch import nn +from transformers import PretrainedConfig + +from sglang.srt.distributed import ( + get_pp_group, + get_tensor_model_parallel_world_size, + parallel_state, + tensor_model_parallel_all_reduce, +) +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.activation import SiluAndMul +from sglang.srt.layers.communicator import ( + LayerCommunicator, + LayerScatterModes, + enable_moe_dense_fully_dp, +) +from sglang.srt.layers.dp_attention import ( + get_attention_dp_size, + get_attention_tp_rank, + get_attention_tp_size, + is_dp_attention_enabled, +) +from sglang.srt.layers.layernorm import RMSNorm +from sglang.srt.layers.linear import ( + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.moe import get_deepep_mode, get_moe_a2a_backend +from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class +from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE +from sglang.srt.layers.moe.token_dispatcher import DeepEPDispatcher +from sglang.srt.layers.moe.topk import TopK +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.radix_attention import AttentionType, RadixAttention +from sglang.srt.layers.rotary_embedding import get_rope +from sglang.srt.layers.utils import PPMissingLayer +from sglang.srt.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode +from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors +from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.models.utils import ( + create_fused_set_kv_buffer_arg, + enable_fused_set_kv_buffer, +) +from sglang.srt.server_args import get_global_server_args +from sglang.srt.utils import add_prefix, is_cuda, is_non_idle_and_non_empty, make_layers + +LoraConfig = None +logger = logging.getLogger(__name__) +_is_cuda = is_cuda() + + +class LLaDA2MoeMLP(nn.Module): + def __init__( + self, + intermediate_size: int, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + reduce_results: Optional[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( + config.hidden_size, + [intermediate_size] * 2, + bias=config.use_bias, + 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, + config.hidden_size, + bias=config.use_bias, + reduce_results=reduce_results, + quant_config=quant_config, + prefix=add_prefix("down_proj", prefix), + tp_rank=tp_rank, + tp_size=tp_size, + ) + + if config.hidden_act != "silu": + raise ValueError("Unsupported activation. Only silu is supported for now.") + self.act_fn = SiluAndMul() + + def forward( + self, + hidden_states: torch.Tensor, + forward_batch: Optional[ForwardBatch] = None, + use_reduce_scatter: bool = False, + ) -> torch.Tensor: + if (self.tp_size == 1) and hidden_states.shape[0] == 0: + return hidden_states + + gate_up, _ = self.gate_up_proj(hidden_states) + hidden_states = self.act_fn(gate_up) + hidden_states, _ = self.down_proj( + hidden_states, skip_all_reduce=use_reduce_scatter + ) + return hidden_states + + +class LLaDA2MoeGate(nn.Module): + def __init__( + self, + config, + params_dtype: Optional[torch.dtype] = None, + prefix: str = "", + ): + super().__init__() + if params_dtype is None: + params_dtype = torch.get_default_dtype() + self.params_dtype = params_dtype + self.weight = nn.Parameter( + torch.empty( + (config.num_experts, config.hidden_size), + dtype=self.params_dtype, + ), + ) + if getattr(config, "moe_router_enable_expert_bias", False): + self.expert_bias = nn.Parameter( + torch.empty((config.num_experts,), dtype=torch.float32), + ) + else: + self.expert_bias = None + + def forward(self, hidden_states): + logits = F.linear(hidden_states.to(self.weight.dtype), self.weight, None).to( + hidden_states.dtype + ) + return logits + + +class LLaDA2MoeSparseMoeBlock(nn.Module): + def __init__( + self, + layer_id: int, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + alt_stream: Optional[torch.cuda.Stream] = None, + prefix: str = "", + ): + super().__init__() + self.layer_id = layer_id + self.alt_stream = alt_stream + self.tp_size = get_tensor_model_parallel_world_size() + self.top_k = config.num_experts_per_tok + self.norm_topk_prob = config.norm_topk_prob + self.hidden_size = config.hidden_size + self.num_shared_experts = config.num_shared_experts + self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0) + self.score_function = getattr(config, "score_function", None) + + if config.hidden_act != "silu": + raise ValueError( + f"Unsupported activation: {config.hidden_act}. " + "Only silu is supported for now." + ) + + # Gate always runs at half / full precision for now. + router_dtype = getattr(config, "router_dtype", None) + if router_dtype is None: + self.router_dtype = None + elif router_dtype == "fp32": + self.router_dtype = torch.float32 + else: + self.router_dtype = torch.bfloat16 + + # TODO global_server_args.ep_num_redundant_experts is used for eplb, not supported now + assert get_global_server_args().ep_num_redundant_experts == 0 + # check group topk + self.num_expert_group = getattr(config, "n_group", 0) + self.topk_group = getattr(config, "topk_group", 0) + if self.num_expert_group > 0 or self.topk_group > 0: + assert ( + self.num_expert_group > 0 + and 0 < self.topk_group <= self.num_expert_group + ) + self.use_grouped_topk = True + else: + self.num_expert_group = self.topk_group = None + self.use_grouped_topk = False + + self.num_experts = ( + config.num_experts + get_global_server_args().ep_num_redundant_experts + ) + + self.gate = LLaDA2MoeGate( + config=config, + params_dtype=self.router_dtype, + prefix=add_prefix("gate", prefix), + ) + self.correction_bias = ( + self.gate.expert_bias.data if self.gate.expert_bias is not None else None + ) + + if self.score_function is not None: + assert ( + self.score_function == "softmax" and self.correction_bias is None + ) or ( + self.score_function == "sigmoid" and self.correction_bias is not None + ), "score_function and correction_bias should be in 2 combination (softmax, None) or (sigmoid, not None)" + + self.topk = TopK( + top_k=self.top_k, + renormalize=self.norm_topk_prob, + use_grouped_topk=self.use_grouped_topk, + num_expert_group=self.num_expert_group, + # num_fused_shared_experts=self.num_fused_shared_experts, + topk_group=self.topk_group, + correction_bias=self.correction_bias, + routed_scaling_factor=self.routed_scaling_factor, + ) + + self.experts = get_moe_impl_class(quant_config)( + num_experts=self.num_experts, + top_k=self.top_k, + layer_id=self.layer_id, + hidden_size=config.hidden_size, + intermediate_size=config.moe_intermediate_size, + quant_config=quant_config, + routed_scaling_factor=self.routed_scaling_factor, + prefix=add_prefix("experts", prefix), + ) + # shared expert + if config.num_shared_experts is not None: + if hasattr(config, "moe_shared_expert_intermediate_size"): + intermediate_size = config.moe_shared_expert_intermediate_size + else: + intermediate_size = config.moe_intermediate_size + intermediate_size *= config.num_shared_experts + # disable tp for shared experts when enable deepep moe + self.shared_experts = LLaDA2MoeMLP( + intermediate_size=intermediate_size, + config=config, + 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() + else {} + ), + ) + # dispatcher + if get_moe_a2a_backend().is_deepep(): + # TODO: we will support tp < ep in the future + self.ep_size = get_tensor_model_parallel_world_size() + + self.deepep_dispatcher = DeepEPDispatcher( + 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.num_experts // self.tp_size, + hidden_size=config.hidden_size, + params_dtype=config.torch_dtype, + deepep_mode=get_deepep_mode(), + async_finish=True, # TODO + return_recv_hook=True, + ) + + def forward( + self, + hidden_states: torch.Tensor, + forward_batch: Optional[ForwardBatch] = None, + use_reduce_scatter: bool = False, + ) -> torch.Tensor: + if not get_moe_a2a_backend().is_deepep(): + return self.forward_normal(hidden_states, use_reduce_scatter) + else: + return self.forward_deepep(hidden_states, forward_batch) + + def get_moe_weights(self): + return [ + x.data + for name, x in self.experts.named_parameters() + if name not in ["correction_bias"] + ] + + def _forward_shared_experts(self, hidden_states: torch.Tensor): + shared_output = None + if self.num_shared_experts > 0: + shared_output = self.shared_experts(hidden_states) + return shared_output + + def _forward_router_experts(self, hidden_states: torch.Tensor): + # router_logits: (num_tokens, n_experts) + router_logits = self.gate(hidden_states) + topk_output = self.topk(hidden_states, router_logits) + return self.experts(hidden_states, topk_output) + + def forward_normal_dual_stream( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + current_stream = torch.cuda.current_stream() + self.alt_stream.wait_stream(current_stream) + shared_output = self._forward_shared_experts(hidden_states.clone()) + + with torch.cuda.stream(self.alt_stream): + router_output = self._forward_router_experts(hidden_states) + current_stream.wait_stream(self.alt_stream) + + return router_output, shared_output + + def forward_normal( + self, + hidden_states: torch.Tensor, + use_reduce_scatter: bool = False, + ) -> torch.Tensor: + num_tokens, hidden_size = hidden_states.shape + hidden_states = hidden_states.view(-1, hidden_size) + + DUAL_STREAM_TOKEN_THRESHOLD = 1024 + if ( + self.alt_stream is not None + and hidden_states.shape[0] > 0 + and hidden_states.shape[0] <= DUAL_STREAM_TOKEN_THRESHOLD + and get_is_capture_mode() + ): + final_hidden_states, shared_output = self.forward_normal_dual_stream( + hidden_states + ) + else: + shared_output = self._forward_shared_experts(hidden_states) + final_hidden_states = self._forward_router_experts(hidden_states) + + if self.num_shared_experts > 0: + final_hidden_states = final_hidden_states + shared_output + + if self.tp_size > 1 and not use_reduce_scatter: + final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) + return final_hidden_states.view(num_tokens, hidden_size) + + def forward_deepep( + self, hidden_states: torch.Tensor, forward_batch: ForwardBatch + ) -> torch.Tensor: + shared_output = None + forward_mode = forward_batch.forward_mode + if is_non_idle_and_non_empty(forward_mode, hidden_states): + router_logits = self.gate(hidden_states) + if self.num_shared_experts > 0: + shared_output = self.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) + + final_hidden_states = self.experts( + hidden_states=hidden_states, + topk_output=topk_output, + ) + + if shared_output is not None: + final_hidden_states += shared_output + return final_hidden_states + + +class LLaDA2MoeAttention(nn.Module): + def __init__( + self, + config: PretrainedConfig, + layer_id: int = 0, + quant_config: Optional[QuantizationConfig] = None, + reduce_results: bool = True, + prefix: str = "", + alt_stream: Optional[torch.cuda.Stream] = None, + ): + super().__init__() + self.hidden_size = config.hidden_size + self.total_num_heads = config.num_attention_heads + self.total_kv_heads = config.num_key_value_heads + self.dp_size = get_attention_dp_size() + attn_tp_rank = get_attention_tp_rank() + attn_tp_size = get_attention_tp_size() + + assert self.total_num_heads % attn_tp_size == 0 + if self.total_kv_heads >= attn_tp_size: + # Number of KV heads is greater than TP size, so we partition + # the KV heads across multiple tensor parallel GPUs. + assert self.total_kv_heads % attn_tp_size == 0 + else: + # Number of KV heads is less than TP size, so we replicate + # the KV heads across multiple tensor parallel GPUs. + assert attn_tp_size % self.total_kv_heads == 0 + assert self.total_num_heads >= self.total_kv_heads + + self.num_heads = self.total_num_heads // attn_tp_size + self.head_dim = config.head_dim or (self.hidden_size // self.total_num_heads) + self.q_size = self.head_dim * self.num_heads + + self.num_kv_heads = max(1, self.total_kv_heads // attn_tp_size) + self.kv_size = max(1, self.num_kv_heads * self.head_dim) + + self.scale = self.head_dim**-0.5 + + self.use_qk_norm = getattr(config, "use_qk_norm", True) + + self.query_key_value = QKVParallelLinear( + self.hidden_size, + self.head_dim, + self.total_num_heads, + self.total_kv_heads, + bias=(config.use_bias or config.use_qkv_bias), + quant_config=quant_config, + prefix=add_prefix("query_key_value", prefix), + tp_rank=attn_tp_rank, + tp_size=attn_tp_size, + ) + + if self.use_qk_norm: + self.query_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.key_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + + self.dense = RowParallelLinear( + self.total_num_heads * self.head_dim, + self.hidden_size, + bias=config.use_bias, + quant_config=quant_config, + reduce_results=reduce_results, + prefix=add_prefix("dense", prefix), + tp_rank=attn_tp_rank, + tp_size=attn_tp_size, + ) + + if hasattr(config, "partial_rotary_factor"): + self.rotary_dim = int(self.head_dim * config.partial_rotary_factor) + elif hasattr(config, "rotary_dim"): + self.rotary_dim = config.rotary_dim + else: + self.rotary_dim = self.head_dim + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.rotary_dim, + max_position=config.max_position_embeddings, + base=config.rope_theta, + rope_scaling=config.rope_scaling, + ) + + self.attn = RadixAttention( + self.num_heads, + self.head_dim, + self.scale, + num_kv_heads=self.num_kv_heads, + layer_id=layer_id, + attn_type=AttentionType.ENCODER_ONLY, + prefix=add_prefix("attn", prefix), + ) + + self.alt_stream = alt_stream + + def _apply_qk_norm( + self, q: torch.Tensor, k: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + # 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_by_head = q.reshape(-1, self.head_dim) + q_by_head = self.query_layernorm(q_by_head) + with torch.cuda.stream(self.alt_stream): + k_by_head = k.reshape(-1, self.head_dim) + k_by_head = self.key_layernorm(k_by_head) + current_stream.wait_stream(self.alt_stream) + else: + q_by_head = q.reshape(-1, self.head_dim) + q_by_head = self.query_layernorm(q_by_head) + k_by_head = k.reshape(-1, self.head_dim) + k_by_head = self.key_layernorm(k_by_head) + q = q_by_head.view(q.shape) + k = k_by_head.view(k.shape) + return q, k + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + if hidden_states.shape[0] == 0: + return hidden_states + qkv, _ = self.query_key_value(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + if self.use_qk_norm: + q, k = self._apply_qk_norm(q, k) + q, k = self.rotary_emb( + positions, + q, + k, + fused_set_kv_buffer_arg=( + create_fused_set_kv_buffer_arg( + value=v, + layer=self.attn, + forward_batch=forward_batch, + ) + if enable_fused_set_kv_buffer(forward_batch) + else None + ), + ) + context_layer = self.attn( + q, + k, + v, + forward_batch, + save_kv_cache=not enable_fused_set_kv_buffer(forward_batch), + ) + attn_output, _ = self.dense(context_layer) + return attn_output + + +class LLaDA2MoeBlock(nn.Module): + def __init__( + self, + config: PretrainedConfig, + layer_id: int = 0, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + alt_stream: Optional[torch.cuda.Stream] = None, + ): + super().__init__() + hidden_size = config.hidden_size + + self.input_layernorm = RMSNorm(hidden_size, eps=config.rms_norm_eps) + self.dp_size = get_attention_dp_size() + self.attention = LLaDA2MoeAttention( + config, + layer_id, + quant_config, + reduce_results=False, + prefix=add_prefix("attention", prefix), + alt_stream=alt_stream, + ) + self.layer_id = layer_id + self.attn_tp_size = get_attention_tp_size() + self.attn_tp_rank = get_attention_tp_rank() + + self.is_layer_sparse = self._is_layer_sparse(config, layer_id=layer_id) + is_previous_layer_sparse = self._is_layer_sparse(config, layer_id=layer_id - 1) + + self.layer_scatter_modes = LayerScatterModes.init_new( + layer_id=layer_id, + num_layers=config.num_hidden_layers, + is_layer_sparse=self.is_layer_sparse, + is_previous_layer_sparse=is_previous_layer_sparse, + ) + + self.is_last_layer = self.layer_id == config.num_hidden_layers - 1 + + if self.is_layer_sparse: + self.mlp = LLaDA2MoeSparseMoeBlock( + layer_id=layer_id, + config=config, + quant_config=quant_config, + alt_stream=alt_stream, + prefix=add_prefix("mlp", prefix), + ) + 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 = LLaDA2MoeMLP( + intermediate_size=config.intermediate_size, + config=config, + quant_config=quant_config, + prefix=add_prefix("mlp", prefix), + tp_rank=mlp_tp_rank, + tp_size=mlp_tp_size, + ) + + self.post_attention_layernorm = RMSNorm(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, + ) + + def _is_layer_sparse(self, config: PretrainedConfig, layer_id: int) -> bool: + return ( + config.num_experts is not None and layer_id >= config.first_k_dense_replace + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + residual: Optional[torch.Tensor], + ) -> torch.Tensor: + hidden_states, residual = self.layer_communicator.prepare_attn( + hidden_states=hidden_states, + residual=residual, + forward_batch=forward_batch, + ) + + hidden_states = self.attention( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + ) + + hidden_states, residual = self.layer_communicator.prepare_mlp( + hidden_states=hidden_states, + residual=residual, + forward_batch=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, use_reduce_scatter) + + hidden_states, residual = self.layer_communicator.postprocess_layer( + hidden_states=hidden_states, + residual=residual, + forward_batch=forward_batch, + ) + + return hidden_states, residual + + +class LLaDA2MoeModel(nn.Module): + + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + alt_stream: Optional[torch.cuda.Stream] = None, + prefix: str = "", + ): + super().__init__() + self.pp_group = get_pp_group() + self.config = config + self.vocab_size = config.vocab_size + self.embed_dim = config.hidden_size + if self.pp_group.is_first_rank: + self.word_embeddings = VocabParallelEmbedding( + self.vocab_size, + self.embed_dim, + quant_config=quant_config, + prefix=add_prefix("word_embeddings", prefix), + enable_tp=not is_dp_attention_enabled(), + ) + else: + self.word_embeddings = PPMissingLayer() + + self.embedding_dropout = torch.nn.Dropout(config.embedding_dropout) + + self.layers, self.start_layer, self.end_layer = make_layers( + config.num_hidden_layers, + lambda idx, prefix: LLaDA2MoeBlock( + layer_id=idx, + config=config, + quant_config=quant_config, + prefix=prefix, + alt_stream=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(self.embed_dim, eps=config.rms_norm_eps) + else: + self.norm = PPMissingLayer(return_tuple=True) + + 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]: + if self.pp_group.is_first_rank: + if input_embeds is None: + hidden_states = self.word_embeddings(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"] + + for i in range(self.start_layer, self.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, + ) + 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 LLaDA2MoeModelLM(nn.Module): + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.pp_group = get_pp_group() + self.config = config + self.quant_config = quant_config + alt_stream = torch.cuda.Stream() if _is_cuda else None + + self.model = LLaDA2MoeModel( + config, + quant_config, + alt_stream=alt_stream, + prefix=add_prefix("model", ""), + ) + + if config.tie_word_embeddings: + self.lm_head = self.model.word_embeddings + else: + # TODO something wrong with ParallelLMHead with DP attention enabled + 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_global_server_args().enable_dp_lm_head, + ) + self.logits_processor = LogitsProcessor(config, return_full_logits=True) + + @property + def start_layer(self): + return self.model.start_layer + + @property + def end_layer(self): + return self.model.end_layer + + def get_embed_and_head(self): + """Used by the eagle_worker.""" + return self.model.word_embeddings.weight, self.lm_head.weight + + def set_embed_and_head(self, embed, head): + """Used by the eagle_worker.""" + del self.model.word_embeddings.weight + del self.lm_head.weight + self.model.word_embeddings.weight = embed + self.lm_head.weight = head + torch.cuda.empty_cache() + torch.cuda.synchronize() + + @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: + hidden_states = self.model( + input_ids, + positions, + forward_batch, + input_embeds, + pp_proxy_tensors=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 + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + + # 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.num_experts, + ) + + params_dict = dict(self.named_parameters()) + for name, loaded_weight in weights: + if ( + ("v_head" in name) + or ("inv_freq" in name) + or (self.config.tie_word_embeddings and "lm_head" in name) + ): + continue + + if ( + hasattr(self.config, "norm_head") + and self.config.norm_head + and "lm_head.weight" in name + ): + import torch.nn.functional as F + + loaded_weight = F.normalize(loaded_weight, dim=0, p=2, eps=1e-7) + + for param_name, weight_name, shard_id in stacked_params_mapping: + 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: + 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 + if name not in params_dict: + continue + + param = params_dict[name] + weight_loader = param.weight_loader + 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) + if name not in params_dict: + continue + param = params_dict[name] + weight_loader = param.weight_loader + 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 + if name not in params_dict: + continue + + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + + self.routed_experts_weights_of_layer = { + layer_id: layer.mlp.get_moe_weights() + for layer_id, layer in enumerate(self.model.layers) + if not isinstance(layer, PPMissingLayer) + and isinstance(layer.mlp, LLaDA2MoeSparseMoeBlock) + } + + @classmethod + def get_model_config_for_expert_location(cls, config): + num_groups = getattr(config, "n_group", 0) + return ModelConfigForExpertLocation( + num_layers=config.num_hidden_layers, + num_logical_experts=config.num_experts, + num_groups=None if num_groups == 0 else num_groups, + ) + + +EntryClass = LLaDA2MoeModelLM diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 9df8ee882..7a4f53202 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -240,6 +240,10 @@ class ServerArgs: revision: Optional[str] = None model_impl: str = "auto" + # Diffusion LLM + dllm_algorithm: Optional[str] = None + dllm_block_size: Optional[int] = None + # HTTP server host: str = "127.0.0.1" port: int = 30000 @@ -663,6 +667,9 @@ class ServerArgs: # Handle exporting request-level metrics. self._handle_request_metrics_exporters() + # Handle diffusion LLM inference. + self._handle_dllm_inference() + # Handle any other necessary validations. self._handle_other_validations() @@ -1974,6 +1981,30 @@ class ServerArgs: "--export-metrics-to-file-dir is required when --export-metrics-to-file is enabled" ) + def _handle_dllm_inference(self): + if self.dllm_algorithm is None: + return + if not self.disable_cuda_graph: + logger.warning( + "Cuda graph is disabled because of using diffusion LLM inference" + ) + self.disable_cuda_graph = True + if not self.disable_overlap_schedule: + logger.warning( + "Overlap schedule is disabled because of using diffusion LLM inference" + ) + self.disable_overlap_schedule = True + if not self.disable_radix_cache: + logger.warning( + "Radix cache is disabled because of using diffusion LLM inference" + ) + self.disable_radix_cache = True + if not self.pp_size > 1: + logger.warning( + "Pipeline parallelism is disabled because of using diffusion LLM inference" + ) + self.pp_size = 1 + def _handle_other_validations(self): # Handle model inference tensor dump. if self.debug_tensor_dump_output_folder is not None: @@ -2093,6 +2124,20 @@ class ServerArgs: "implementation.\n", ) + # Diffusion LLM + parser.add_argument( + "--dllm-algorithm", + type=str, + default=ServerArgs.dllm_algorithm, + help="The diffusion LLM algorithm.", + ) + parser.add_argument( + "--dllm-block-size", + type=int, + default=ServerArgs.dllm_block_size, + help="The number of tokens processed in each iteration of the block diffusion LLM.", + ) + # HTTP server parser.add_argument( "--host",