sglang 0.5.3rc0__py3-none-any.whl → 0.5.3rc2__py3-none-any.whl

sglang/srt/layers/attention/trtllm_mla_backend.py

@@ -22,7 +22,7 @@ from sglang.srt.layers.attention.utils import (
 from sglang.srt.layers.dp_attention import get_attention_tp_size
 from sglang.srt.managers.schedule_batch import global_server_args_dict
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
-from sglang.srt.utils import is_flashinfer_available
+from sglang.srt.utils import is_cuda, is_flashinfer_available
 
 if is_flashinfer_available():
     import flashinfer
@@ -30,7 +30,12 @@ if is_flashinfer_available():
 if TYPE_CHECKING:
     from sglang.srt.layers.radix_attention import RadixAttention
     from sglang.srt.model_executor.model_runner import ModelRunner
-    from sglang.srt.speculative.spec_info import SpecInfo
+    from sglang.srt.speculative.spec_info import SpecInput
+
+_is_cuda = is_cuda()
+
+if _is_cuda:
+    from sgl_kernel import concat_mla_absorb_q
 
 # Constants
 DEFAULT_WORKSPACE_SIZE_MB = 128  # Memory workspace size in MB
@@ -122,6 +127,8 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
             "disable_chunked_prefix_cache"
         ]
 
+        self.num_draft_tokens = model_runner.server_args.speculative_num_draft_tokens
+
     def _calc_padded_blocks(self, max_seq_len: int) -> int:
         """
         Calculate padded block count that satisfies both TRT-LLM and Triton constraints.
@@ -207,12 +214,12 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
         seq_lens: torch.Tensor,
         encoder_lens: Optional[torch.Tensor],
         forward_mode: ForwardMode,
-        spec_info: Optional[SpecInfo],
+        spec_info: Optional[SpecInput],
     ):
         """Initialize metadata for CUDA graph capture."""
 
         # Delegate to parent for non-decode modes.
-        if not forward_mode.is_decode_or_idle():
+        if not forward_mode.is_decode_or_idle() and not forward_mode.is_target_verify():
             return super().init_forward_metadata_capture_cuda_graph(
                 bs,
                 num_tokens,
@@ -223,6 +230,9 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
                 spec_info,
             )
 
+        if forward_mode.is_target_verify():
+            seq_lens = seq_lens + self.num_draft_tokens
+
         # Custom fast-path for decode/idle.
         # Capture with full width so future longer sequences are safe during replay
         max_blocks_per_seq = self._calc_padded_blocks(self.max_context_len)
@@ -260,12 +270,12 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
         seq_lens_sum: int,
         encoder_lens: Optional[torch.Tensor],
         forward_mode: ForwardMode,
-        spec_info: Optional[SpecInfo],
+        spec_info: Optional[SpecInput],
         seq_lens_cpu: Optional[torch.Tensor],
     ):
         """Replay CUDA graph with new inputs."""
         # Delegate to parent for non-decode modes.
-        if not forward_mode.is_decode_or_idle():
+        if not forward_mode.is_decode_or_idle() and not forward_mode.is_target_verify():
             return super().init_forward_metadata_replay_cuda_graph(
                 bs,
                 req_pool_indices,
@@ -277,6 +287,10 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
                 seq_lens_cpu,
             )
 
+        if forward_mode.is_target_verify():
+            seq_lens = seq_lens + self.num_draft_tokens
+            del seq_lens_sum  # not handle "num_draft_tokens" but we do not need it
+
         metadata = self.decode_cuda_graph_metadata[bs]
 
         # Update block indices for new sequences.
@@ -327,7 +341,10 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
                 cum_seq_lens_q,
                 seq_lens,
             )
-        elif forward_batch.forward_mode.is_decode_or_idle():
+        elif (
+            forward_batch.forward_mode.is_decode_or_idle()
+            or forward_batch.forward_mode.is_target_verify()
+        ):
             bs = forward_batch.batch_size
 
             # Get maximum sequence length.
@@ -336,13 +353,19 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
             else:
                 max_seq = forward_batch.seq_lens.max().item()
 
+            seq_lens = forward_batch.seq_lens
+
+            if forward_batch.forward_mode.is_target_verify():
+                max_seq = max_seq + self.num_draft_tokens
+                seq_lens = seq_lens + self.num_draft_tokens
+
             max_seqlen_pad = self._calc_padded_blocks(max_seq)
             block_kv_indices = self._create_block_kv_indices(
                 bs,
                 max_seqlen_pad,
                 forward_batch.req_pool_indices,
-                forward_batch.seq_lens,
-                forward_batch.seq_lens.device,
+                seq_lens,
+                seq_lens.device,
             )
 
             max_seq_len_val = int(max_seq)
@@ -482,7 +505,7 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
             q_rope_reshaped = q_rope.view(
                 -1, layer.tp_q_head_num, layer.head_dim - layer.v_head_dim
             )
-            query = torch.cat([q_nope, q_rope_reshaped], dim=-1)
+            query = _concat_mla_absorb_q_general(q_nope, q_rope_reshaped)
         else:
             # For FP8 path, we already have the query and rope parts merged because of the quantize_and_rope_for_fp8 function
             query = q.view(-1, layer.tp_q_head_num, layer.head_dim)
@@ -545,49 +568,163 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
         save_kv_cache: bool = True,
         q_rope: Optional[torch.Tensor] = None,
         k_rope: Optional[torch.Tensor] = None,
+        cos_sin_cache: Optional[torch.Tensor] = None,
+        is_neox: Optional[bool] = False,
     ) -> torch.Tensor:
-        if (
-            forward_batch.forward_mode.is_target_verify()
-            or forward_batch.forward_mode.is_draft_extend()
-        ):
+        if forward_batch.forward_mode.is_draft_extend():
             return super().forward_extend(
                 q, k, v, layer, forward_batch, save_kv_cache, q_rope, k_rope
             )
-        # chunked prefix cache is not enabled, use Flashinfer MLA prefill kernel
-        if forward_batch.attn_attend_prefix_cache is None:
-            return super().forward_extend(
-                q, k, v, layer, forward_batch, save_kv_cache, q_rope, k_rope
+
+        # TODO refactor to avoid code duplication
+        merge_query = q_rope is not None
+        if (
+            self.data_type == torch.float8_e4m3fn
+        ) and forward_batch.forward_mode.is_target_verify():
+            # For FP8 path, we quantize the query and rope parts and merge them into a single tensor
+            # Note: rope application in deepseek_v2.py:forward_absorb_prepare is skipped for FP8 decode path of this trtllm_mla backend
+            assert all(
+                x is not None for x in [q_rope, k_rope, cos_sin_cache]
+            ), "For FP8 path and using flashinfer.rope.mla_rope_quantize we need all of q_rope, k_rope and cos_sin_cache to be not None."
+            q, k, k_rope = self.quantize_and_rope_for_fp8(
+                q,
+                q_rope,
+                k.squeeze(1),
+                k_rope.squeeze(1),
+                forward_batch,
+                cos_sin_cache,
+                is_neox,
+            )
+            merge_query = False
+
+        # Save KV cache if requested
+        if save_kv_cache:
+            assert (
+                k is not None and k_rope is not None
+            ), "For populating trtllm_mla kv cache, both k_nope and k_rope should be not None."
+            forward_batch.token_to_kv_pool.set_mla_kv_buffer(
+                layer, forward_batch.out_cache_loc, k, k_rope
             )
 
-        if not forward_batch.attn_attend_prefix_cache:
+        # TODO refactor to avoid code duplication
+        # Prepare query tensor inline
+        if merge_query:
+            # For FP16 path, we merge the query and rope parts into a single tensor
+            q_nope = q.view(-1, layer.tp_q_head_num, layer.v_head_dim)
+            q_rope_reshaped = q_rope.view(
+                -1, layer.tp_q_head_num, layer.head_dim - layer.v_head_dim
+            )
+            q = _concat_mla_absorb_q_general(q_nope, q_rope_reshaped)
+        else:
+            # For FP8 path, we already have the query and rope parts merged because of the quantize_and_rope_for_fp8 function
             q = q.view(-1, layer.tp_q_head_num, layer.head_dim)
-            k = k.view(-1, layer.tp_k_head_num, layer.head_dim)
-            v = v.view(-1, layer.tp_k_head_num, layer.v_head_dim)
-            output = flashinfer.prefill.trtllm_ragged_attention_deepseek(
+
+        q = q.view(-1, layer.tp_q_head_num, layer.head_dim)
+
+        if k_rope is not None:
+            k = torch.cat([k, k_rope], dim=-1)
+        k = k.view(-1, layer.tp_k_head_num, layer.head_dim)
+
+        v = v.view(-1, layer.tp_k_head_num, layer.v_head_dim)
+
+        if forward_batch.forward_mode.is_target_verify():
+            metadata = (
+                getattr(forward_batch, "decode_trtllm_mla_metadata", None)
+                or self.forward_decode_metadata
+            )
+
+            # Ensure query has shape [bs, num_draft_tokens, num_q_heads, head_dim]
+            bs = forward_batch.batch_size
+            q = q.view(bs, -1, layer.tp_q_head_num, layer.head_dim)
+
+            k_cache = forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id)
+            kv_cache = k_cache.view(-1, self.page_size, self.kv_cache_dim).unsqueeze(1)
+
+            q_scale = 1.0
+            k_scale = (
+                layer.k_scale_float
+                if getattr(layer, "k_scale_float", None) is not None
+                else 1.0
+            )
+
+            bmm1_scale = q_scale * k_scale * layer.scaling
+
+            seq_lens = (
+                forward_batch.seq_lens.to(torch.int32)
+                + forward_batch.spec_info.draft_token_num
+            )
+            max_seq_len = metadata.max_seq_len + forward_batch.spec_info.draft_token_num
+
+            # TODO may use `mla_rope_quantize_fp8` fusion
+            q = q.to(self.data_type)
+            assert kv_cache.dtype == self.data_type
+
+            raw_out = flashinfer.decode.trtllm_batch_decode_with_kv_cache_mla(
+                query=q,
+                kv_cache=kv_cache,
+                workspace_buffer=self.workspace_buffer,
+                qk_nope_head_dim=self.qk_nope_head_dim,
+                kv_lora_rank=self.kv_lora_rank,
+                qk_rope_head_dim=self.qk_rope_head_dim,
+                block_tables=metadata.block_kv_indices,
+                seq_lens=seq_lens,
+                max_seq_len=max_seq_len,
+                bmm1_scale=bmm1_scale,
+            )
+
+            # Reshape output directly without slicing
+            output = raw_out.view(-1, layer.tp_q_head_num * layer.v_head_dim)
+            return output
+
+        if forward_batch.attn_attend_prefix_cache:
+            # MHA for chunked prefix kv cache when running model with MLA
+            assert forward_batch.prefix_chunk_idx is not None
+            assert forward_batch.prefix_chunk_cu_seq_lens is not None
+            assert q_rope is None
+            assert k_rope is None
+            chunk_idx = forward_batch.prefix_chunk_idx
+
+            output_shape = (q.shape[0], layer.tp_q_head_num, layer.v_head_dim)
+            return flashinfer.prefill.trtllm_ragged_attention_deepseek(
                 query=q,
                 key=k,
                 value=v,
                 workspace_buffer=self.workspace_buffer,
-                seq_lens=self.forward_prefill_metadata.seq_lens,
+                seq_lens=forward_batch.prefix_chunk_seq_lens[chunk_idx],
                 max_q_len=self.forward_prefill_metadata.max_seq_len,
-                max_kv_len=self.forward_prefill_metadata.max_seq_len,
+                max_kv_len=forward_batch.prefix_chunk_max_seq_lens[chunk_idx],
                 bmm1_scale=layer.scaling,
                 bmm2_scale=1.0,
-                o_sf_scale=1.0,
+                o_sf_scale=-1.0,
                 batch_size=forward_batch.batch_size,
                 window_left=-1,
                 cum_seq_lens_q=self.forward_prefill_metadata.cum_seq_lens,
-                cum_seq_lens_kv=self.forward_prefill_metadata.cum_seq_lens,
+                cum_seq_lens_kv=forward_batch.prefix_chunk_cu_seq_lens[chunk_idx],
                 enable_pdl=False,
-                is_causal=True,
-                return_lse=forward_batch.mha_return_lse,
+                is_causal=False,
+                return_lse=True,
+                out=torch.zeros(*output_shape, dtype=q.dtype, device=q.device),
             )
-        else:
-            # replace with trtllm ragged attention once accuracy is resolved.
-            output = super().forward_extend(
-                q, k, v, layer, forward_batch, save_kv_cache, q_rope, k_rope
-            )
-        return output
+
+        return flashinfer.prefill.trtllm_ragged_attention_deepseek(
+            query=q,
+            key=k,
+            value=v,
+            workspace_buffer=self.workspace_buffer,
+            seq_lens=self.forward_prefill_metadata.seq_lens,
+            max_q_len=self.forward_prefill_metadata.max_seq_len,
+            max_kv_len=self.forward_prefill_metadata.max_seq_len,
+            bmm1_scale=layer.scaling,
+            bmm2_scale=1.0,
+            o_sf_scale=1.0,
+            batch_size=forward_batch.batch_size,
+            window_left=-1,
+            cum_seq_lens_q=self.forward_prefill_metadata.cum_seq_lens,
+            cum_seq_lens_kv=self.forward_prefill_metadata.cum_seq_lens,
+            enable_pdl=False,
+            is_causal=True,
+            return_lse=forward_batch.mha_return_lse,
+        )
 
 
 class TRTLLMMLAMultiStepDraftBackend(FlashInferMLAMultiStepDraftBackend):
@@ -605,3 +742,10 @@ class TRTLLMMLAMultiStepDraftBackend(FlashInferMLAMultiStepDraftBackend):
                 kv_indptr_buf=self.kv_indptr[i],
                 q_indptr_decode_buf=self.q_indptr_decode,
             )
+
+
+def _concat_mla_absorb_q_general(q_nope, q_rope):
+    if _is_cuda and q_nope.shape[-1] == 512 and q_rope.shape[-1] == 64:
+        return concat_mla_absorb_q(q_nope, q_rope)
+    else:
+        return torch.cat([q_nope, q_rope], dim=-1)

sglang/srt/layers/attention/vision.py

@@ -16,14 +16,19 @@ from sglang.srt.utils import (
     get_device_capability,
     is_blackwell,
     is_cuda,
+    is_npu,
     print_info_once,
 )
 
 _is_cuda = is_cuda()
+_is_npu = is_npu()
 
 if _is_cuda:
     from sgl_kernel.flash_attn import flash_attn_varlen_func
 
+if _is_npu:
+    import torch_npu
+
 from sglang.srt.distributed import (
     split_tensor_along_last_dim,
     tensor_model_parallel_all_gather,
@@ -331,10 +336,63 @@ class VisionFlash3Attention(nn.Module):
         return output
 
 
+class VisionAscendAttention(nn.Module):
+
+    def __init__(
+        self,
+        **kwargs,
+    ):
+        if not _is_npu:
+            raise Exception("VisionAscendAttention is only available for ascend npu")
+        super().__init__()
+
+    def forward(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        cu_seqlens: Optional[Union[SingletonCache, torch.Tensor]],
+        bsz: int,
+        seq_len: int,
+        **kwargs,
+    ) -> torch.Tensor:
+        r"""
+        Args:
+            cu_seqlens: [b]
+        Returns:
+             [b * s, h, head_size]
+        """
+        if cu_seqlens is None:
+            cu_seqlens = _get_cu_seqlens_for_shape(bsz, seq_len, device=q.device)
+
+        seq_lens = cu_seqlens[1:] - cu_seqlens[:-1]
+        if seq_lens.is_npu:
+            # cu_seqlens must be on cpu because of operator restriction
+            seq_lens = seq_lens.to("cpu")
+        _, num_heads, head_size = q.shape
+        num_kv_heads = k.shape[1]
+        output = torch.empty_like(q)
+
+        # operator requires pta version >= 2.5.1
+        torch_npu._npu_flash_attention_unpad(
+            query=q,
+            key=k,
+            value=v,
+            seq_len=seq_lens.to(torch.int32),
+            scale_value=head_size**-0.5,
+            num_heads=num_heads,
+            num_kv_heads=num_kv_heads,
+            out=output,
+        )
+
+        return output
+
+
 QKV_BACKEND_IMPL = {
     "triton_attn": VisionTritonAttention,
     "sdpa": VisionSdpaAttention,
     "fa3": VisionFlash3Attention,
+    "ascend_attn": VisionAscendAttention,
 }
 
 

sglang/srt/layers/attention/wave_backend.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import logging
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Optional, Union
+from typing import TYPE_CHECKING, Optional
 
 import torch
 import triton
@@ -17,7 +17,7 @@ from sglang.srt.utils import get_bool_env_var, get_device_core_count
 if TYPE_CHECKING:
     from sglang.srt.layers.radix_attention import RadixAttention
     from sglang.srt.model_executor.model_runner import ModelRunner
-    from sglang.srt.speculative.eagle_utils import EagleDraftInput, EagleVerifyInput
+    from sglang.srt.speculative.spec_info import SpecInput
 
 logger = logging.getLogger(__name__)
 
@@ -393,7 +393,7 @@ class WaveAttnBackend(AttentionBackend):
         seq_lens: torch.Tensor,
         encoder_lens: Optional[torch.Tensor],
         forward_mode: ForwardMode,
-        spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]],
+        spec_info: Optional[SpecInput],
     ):
         assert encoder_lens is None, "Not supported"
 
@@ -477,7 +477,7 @@ class WaveAttnBackend(AttentionBackend):
         seq_lens_sum: int,
         encoder_lens: Optional[torch.Tensor],
         forward_mode: ForwardMode,
-        spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]],
+        spec_info: Optional[SpecInput],
         seq_lens_cpu: Optional[torch.Tensor],
     ):
         # NOTE: encoder_lens expected to be zeros or None

sglang/srt/layers/communicator.py

@@ -50,6 +50,7 @@ from sglang.srt.utils import (
     is_hip,
     is_sm90_supported,
     is_sm100_supported,
+    prepare_weight_cache,
 )
 
 _is_flashinfer_available = is_flashinfer_available()
@@ -275,7 +276,11 @@ class LayerCommunicator:
         hidden_states: torch.Tensor,
         residual: torch.Tensor,
         forward_batch: ForwardBatch,
+        cache=None,
     ):
+        if cache is not None:
+            self._context.cache = cache
+
         return self._communicate_with_all_reduce_and_layer_norm_fn(
             hidden_states=hidden_states,
             residual=residual,
@@ -349,6 +354,7 @@ class CommunicateContext:
     attn_tp_size: int
     attn_dp_size: int
     tp_size: int
+    cache = None
 
     def is_same_group_size(self, a: ScatterMode, b: ScatterMode):
         return self.process_group_sizes[a] == self.process_group_sizes[b]
@@ -533,6 +539,8 @@ class CommunicateWithAllReduceAndLayerNormFn:
                 )
             else:
                 hidden_states = tensor_model_parallel_all_reduce(hidden_states)
+                if context.cache is not None:
+                    _ = prepare_weight_cache(hidden_states, context.cache)
                 hidden_states, residual = layernorm(hidden_states, residual)
         return hidden_states, residual
 

sglang/srt/layers/dp_attention.py

@@ -17,6 +17,7 @@ from sglang.srt.distributed import (
     get_tp_group,
     tensor_model_parallel_all_reduce,
 )
+from sglang.srt.utils import get_bool_env_var, is_hip
 
 if TYPE_CHECKING:
     from sglang.srt.configs.model_config import ModelConfig
@@ -36,6 +37,9 @@ _LOCAL_ATTN_DP_SIZE: Optional[int] = None
 _LOCAL_ATTN_DP_RANK: Optional[int] = None
 _ENABLE_DP_ATTENTION_FLAG: bool = False
 
+_is_hip = is_hip()
+_USE_ROCM700A_WA = _is_hip and get_bool_env_var("SGLANG_USE_ROCM700A")
+
 
 class DpPaddingMode(IntEnum):
 
@@ -67,7 +71,12 @@ class DpPaddingMode(IntEnum):
 
     @classmethod
     def get_default_mode_in_cuda_graph(cls) -> DpPaddingMode:
-        return cls.MAX_LEN
+        # TODO(kkhuang-amd): noqa, temporary work-around for rocm 7.0.0 alpha
+        # it can be safely removed later, once RCCL fixed
+        if _USE_ROCM700A_WA:
+            return cls.SUM_LEN
+        else:
+            return cls.MAX_LEN
 
 
 class _DpGatheredBufferWrapper:
@@ -254,6 +263,7 @@ def initialize_dp_attention(
         use_pynccl=SYNC_TOKEN_IDS_ACROSS_TP,
         use_pymscclpp=False,
         use_custom_allreduce=False,
+        use_torch_symm_mem=False,
         use_hpu_communicator=False,
         use_xpu_communicator=False,
         use_npu_communicator=False,

sglang/srt/layers/elementwise.py

@@ -187,7 +187,9 @@ fused_dual_residual_rmsnorm_kernel_autotune = rmsnorm_autotune(
 
 def fused_dual_residual_rmsnorm(x, residual, weight1, weight2, eps, autotune=False):
     assert len(x.shape) == 2
-    assert x.shape == residual.shape and x.dtype == residual.dtype
+    assert (
+        x.shape == residual.shape and x.dtype == residual.dtype
+    ), f"{x.shape=} {residual.shape=} {x.dtype=} {residual.dtype=}"
     output, mid = torch.empty_like(x), torch.empty_like(x)
     bs, hidden_dim = x.shape
     if autotune:

sglang/srt/layers/layernorm.py

@@ -80,6 +80,8 @@ class RMSNorm(CustomOp):
         )
         if _use_aiter:
             self._forward_method = self.forward_aiter
+        if get_bool_env_var("SGLANG_ENABLE_DETERMINISTIC_INFERENCE"):
+            self._forward_method = self.forward_native
 
     def forward_cuda(
         self,

sglang/srt/layers/linear.py

@@ -31,6 +31,7 @@ from sglang.srt.layers.parameter import (
     _ColumnvLLMParameter,
 )
 from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod
+from sglang.srt.layers.utils import pad_or_narrow_weight
 from sglang.srt.utils import is_cpu, is_npu, set_weight_attrs
 
 if TYPE_CHECKING:
@@ -625,9 +626,16 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
                 # bitsandbytes loads the weights of the specific portion
                 # no need to narrow here
                 if not use_bitsandbytes_4bit and not self.use_presharded_weights:
-                    loaded_weight = loaded_weight.narrow(
-                        output_dim, start_idx, shard_size
-                    )
+                    # Padding for special case like qwen2_5_VL's mlp which is not 8-aligned
+                    end_idx = start_idx + shard_size
+                    if end_idx > loaded_weight.shape[output_dim]:
+                        loaded_weight = pad_or_narrow_weight(
+                            loaded_weight, output_dim, start_idx, shard_size
+                        )
+                    else:
+                        loaded_weight = loaded_weight.narrow(
+                            output_dim, start_idx, shard_size
+                        )
 
         # Special case for AQLM codebooks.
         elif is_metadata:
@@ -1302,7 +1310,16 @@ class RowParallelLinear(LinearBase):
                     shard_size,
                 )
             else:
-                loaded_weight = loaded_weight.narrow(input_dim, start_idx, shard_size)
+                # Padding for special case like qwen2_5_VL's mlp which is not 8-aligned
+                end_idx = start_idx + shard_size
+                if end_idx > loaded_weight.shape[input_dim]:
+                    loaded_weight = pad_or_narrow_weight(
+                        loaded_weight, input_dim, start_idx, shard_size
+                    )
+                else:
+                    loaded_weight = loaded_weight.narrow(
+                        input_dim, start_idx, shard_size
+                    )
 
         # Special case for loading scales off disk, which often do not
         # have a shape (such as in the case of AutoFP8).

sglang/srt/layers/logits_processor.py

@@ -220,6 +220,7 @@ class LogitsProcessor(nn.Module):
         self.config = config
         self.logit_scale = logit_scale
         self.use_attn_tp_group = global_server_args_dict["enable_dp_lm_head"]
+        self.use_fp32_lm_head = global_server_args_dict["enable_fp32_lm_head"]
         if self.use_attn_tp_group:
             self.attn_tp_size = get_attention_tp_size()
             self.do_tensor_parallel_all_gather = (
@@ -461,7 +462,11 @@ class LogitsProcessor(nn.Module):
             dp_gather_replicate(hidden_states, local_hidden_states, logits_metadata)
 
         if hasattr(lm_head, "weight"):
-            if use_intel_amx_backend(lm_head):
+            if self.use_fp32_lm_head:
+                logits = torch.matmul(
+                    hidden_states.to(torch.float32), lm_head.weight.to(torch.float32).T
+                )
+            elif use_intel_amx_backend(lm_head):
                 logits = torch.ops.sgl_kernel.weight_packed_linear(
                     hidden_states.to(lm_head.weight.dtype),
                     lm_head.weight,
@@ -475,7 +480,15 @@ class LogitsProcessor(nn.Module):
         else:
             # GGUF models
             # TODO: use weight_packed_linear for GGUF models
-            logits = lm_head.quant_method.apply(lm_head, hidden_states, embedding_bias)
+            if self.use_fp32_lm_head:
+                with torch.cuda.amp.autocast(enabled=False):
+                    logits = lm_head.quant_method.apply(
+                        lm_head, hidden_states.to(torch.float32), embedding_bias
+                    )
+            else:
+                logits = lm_head.quant_method.apply(
+                    lm_head, hidden_states, embedding_bias
+                )
 
         if self.logit_scale is not None:
             logits.mul_(self.logit_scale)

sglang/srt/layers/moe/ep_moe/kernels.py

@@ -1104,10 +1104,10 @@ def ep_gather(
     input_index: torch.Tensor,
     output_tensor: torch.Tensor,
 ):
-    BLOCK_D = 1024 if not is_in_ci() else 128  # block size of quantization
     num_warps = 2
     num_tokens = output_tensor.shape[0]
     hidden_size = input_tensor.shape[1]
+    BLOCK_D = 128 if hidden_size % 1024 != 0 else 1024  # block size of quantization
     assert hidden_size % BLOCK_D == 0
     grid = (triton.cdiv(hidden_size, BLOCK_D), min(num_tokens, 1024))
     _fwd_kernel_ep_gather[grid](