sglang 0.5.1.post3__py3-none-any.whl → 0.5.2rc1__py3-none-any.whl

sglang/srt/layers/attention/ascend_backend.py

@@ -64,7 +64,7 @@ class AscendAttnBackend(AttentionBackend):
         if self.use_mla:
             self.kv_lora_rank = model_runner.model_config.kv_lora_rank
             self.qk_rope_head_dim = model_runner.model_config.qk_rope_head_dim
-            self.native_attn = TorchNativeAttnBackend(model_runner)
+        self.native_attn = TorchNativeAttnBackend(model_runner)
         self.graph_metadata = {}
         self.max_context_len = model_runner.model_config.context_len
         self.req_to_token = model_runner.req_to_token_pool.req_to_token
@@ -158,7 +158,7 @@ class AscendAttnBackend(AttentionBackend):
         self.graph_mode = True
 
     def get_cuda_graph_seq_len_fill_value(self):
-        return 1
+        return 0
 
     def forward_extend(
         self,
@@ -167,7 +167,7 @@ class AscendAttnBackend(AttentionBackend):
         v,
         layer: RadixAttention,
         forward_batch: ForwardBatch,
-        save_kv_cache=True,
+        save_kv_cache: bool = True,
     ):
         if not self.use_mla:
             if save_kv_cache:
@@ -180,7 +180,7 @@ class AscendAttnBackend(AttentionBackend):
 
             if self.use_fia:
                 """FIA will support multi-bs in the later version of CANN"""
-                q = q.view(-1, layer.tp_q_head_num, layer.qk_head_dim)
+                q = q.reshape(-1, layer.tp_q_head_num, layer.qk_head_dim)
                 attn_output = torch.empty(
                     (q.size(0), layer.tp_q_head_num, layer.v_head_dim),
                     device=q.device,
@@ -208,26 +208,61 @@ class AscendAttnBackend(AttentionBackend):
                 )
 
             else:
-                query = q.view(-1, layer.tp_q_head_num * layer.qk_head_dim)
-                attn_output = torch.empty(
-                    (query.shape[0], layer.tp_q_head_num * layer.v_head_dim),
-                    dtype=query.dtype,
-                    device=query.device,
-                )
+                if layer.qk_head_dim <= 128:
+                    query = q.reshape(-1, layer.tp_q_head_num * layer.qk_head_dim)
+                    attn_output = torch.empty(
+                        (query.shape[0], layer.tp_q_head_num * layer.v_head_dim),
+                        dtype=query.dtype,
+                        device=query.device,
+                    )
 
-                torch_npu._npu_flash_attention_qlens(
-                    query=query,
-                    key_cache=k_cache,
-                    value_cache=v_cache,
-                    mask=self.mask,
-                    block_table=self.forward_metadata.block_tables,
-                    seq_len=self.forward_metadata.extend_seq_lens_cpu_int,
-                    context_lens=self.forward_metadata.seq_lens_cpu_int,
-                    scale_value=layer.scaling,
-                    num_heads=layer.tp_q_head_num,
-                    num_kv_heads=layer.tp_k_head_num,
-                    out=attn_output,
-                )
+                    torch_npu._npu_flash_attention_qlens(
+                        query=query,
+                        key_cache=k_cache,
+                        value_cache=v_cache,
+                        mask=self.mask,
+                        block_table=self.forward_metadata.block_tables,
+                        seq_len=self.forward_metadata.extend_seq_lens_cpu_int,
+                        context_lens=self.forward_metadata.seq_lens_cpu_int,
+                        scale_value=layer.scaling,
+                        num_heads=layer.tp_q_head_num,
+                        num_kv_heads=layer.tp_k_head_num,
+                        out=attn_output,
+                    )
+                else:
+                    if layer.qk_head_dim != layer.v_head_dim:
+                        attn_output = q.new_empty(
+                            (q.shape[0], layer.tp_q_head_num * layer.v_head_dim)
+                        )
+                    else:
+                        attn_output = torch.empty_like(q)
+
+                    use_gqa = layer.tp_q_head_num != layer.tp_k_head_num
+
+                    q_ = q.view(-1, layer.tp_q_head_num, layer.qk_head_dim)
+                    o_ = attn_output.view(-1, layer.tp_q_head_num, layer.v_head_dim)
+
+                    causal = True
+                    if (
+                        layer.is_cross_attention
+                        or layer.attn_type == AttentionType.ENCODER_ONLY
+                    ):
+                        causal = False
+
+                    self.native_attn._run_sdpa_forward_extend(
+                        q_,
+                        o_,
+                        k_cache.view(-1, layer.tp_k_head_num, layer.qk_head_dim),
+                        v_cache.view(-1, layer.tp_v_head_num, layer.v_head_dim),
+                        forward_batch.req_to_token_pool.req_to_token,
+                        forward_batch.req_pool_indices,
+                        forward_batch.seq_lens,
+                        forward_batch.extend_prefix_lens,
+                        forward_batch.extend_seq_lens,
+                        scaling=layer.scaling,
+                        enable_gqa=use_gqa,
+                        causal=causal,
+                    )
         else:
             assert (
                 layer.qk_head_dim != layer.v_head_dim
@@ -253,6 +288,136 @@ class AscendAttnBackend(AttentionBackend):
 
         return attn_output
 
+    def forward_decode_graph(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        layer: RadixAttention,
+        forward_batch: ForwardBatch,
+        save_kv_cache: bool = True,
+        q_rope: Optional[torch.Tensor] = None,
+        k_rope: Optional[torch.Tensor] = None,
+    ):
+        if save_kv_cache:
+            if self.use_mla:
+                k = k.view(-1, layer.tp_k_head_num, self.kv_lora_rank)
+                k_rope = k_rope.view(-1, layer.tp_k_head_num, self.qk_rope_head_dim)
+                forward_batch.token_to_kv_pool.set_kv_buffer(
+                    layer, forward_batch.out_cache_loc, k, k_rope
+                )
+            else:
+                forward_batch.token_to_kv_pool.set_kv_buffer(
+                    layer, forward_batch.out_cache_loc, k, v
+                )
+
+        if not self.use_mla:
+            k_cache = forward_batch.token_to_kv_pool.get_key_buffer(
+                layer.layer_id
+            ).view(-1, self.page_size, layer.tp_k_head_num * layer.qk_head_dim)
+            v_cache = forward_batch.token_to_kv_pool.get_value_buffer(
+                layer.layer_id
+            ).view(-1, self.page_size, layer.tp_v_head_num * layer.v_head_dim)
+            query = q.reshape(-1, 1, layer.tp_q_head_num * layer.qk_head_dim)
+            if self.forward_metadata.seq_lens_cpu_int is None:
+                actual_seq_len_kv = self.forward_metadata.seq_lens_cpu_list
+            else:
+                actual_seq_len_kv = (
+                    self.forward_metadata.seq_lens_cpu_int.cpu().int().tolist()
+                )
+            num_tokens = query.shape[0]
+            workspace = torch_npu._npu_fused_infer_attention_score_get_max_workspace(
+                query,
+                k_cache,
+                v_cache,
+                block_table=self.forward_metadata.block_tables,
+                block_size=self.page_size,
+                num_heads=layer.tp_q_head_num,
+                num_key_value_heads=layer.tp_k_head_num,
+                input_layout="BSH",
+                scale=layer.scaling,
+                actual_seq_lengths_kv=actual_seq_len_kv,
+            )
+            output = torch.empty(
+                (num_tokens, 1, layer.tp_q_head_num * layer.v_head_dim),
+                dtype=q.dtype,
+                device=q.device,
+            )
+            softmax_lse = torch.empty(1, dtype=q.dtype, device=q.device)
+            torch_npu.npu_fused_infer_attention_score.out(
+                query,
+                k_cache,
+                v_cache,
+                block_table=self.forward_metadata.block_tables,
+                block_size=self.page_size,
+                num_heads=layer.tp_q_head_num,
+                num_key_value_heads=layer.tp_k_head_num,
+                input_layout="BSH",
+                scale=layer.scaling,
+                actual_seq_lengths_kv=actual_seq_len_kv,
+                workspace=workspace,
+                out=[output, softmax_lse],
+            )
+            return output.view(num_tokens, layer.tp_q_head_num * layer.v_head_dim)
+        else:
+            c_kv, k_rope = forward_batch.token_to_kv_pool.get_kv_buffer(layer.layer_id)
+            k_rope_cache = k_rope.view(
+                -1, layer.tp_k_head_num, self.page_size, self.qk_rope_head_dim
+            )
+            c_kv_cache = c_kv.view(
+                -1, layer.tp_v_head_num, self.page_size, self.kv_lora_rank
+            )
+
+            q_nope = q.view(-1, layer.tp_q_head_num, 1, self.kv_lora_rank)
+            q_rope = q_rope.view(-1, layer.tp_q_head_num, 1, self.qk_rope_head_dim)
+            if self.forward_metadata.seq_lens_cpu_int is None:
+                actual_seq_len_kv = self.forward_metadata.seq_lens_cpu_list
+            else:
+                actual_seq_len_kv = (
+                    self.forward_metadata.seq_lens_cpu_int.cpu().int().tolist()
+                )
+
+            workspace = torch_npu._npu_fused_infer_attention_score_get_max_workspace(
+                q_nope,
+                c_kv_cache,
+                c_kv_cache,
+                query_rope=q_rope,
+                key_rope=k_rope_cache,
+                num_heads=layer.tp_q_head_num,
+                num_key_value_heads=layer.tp_k_head_num,
+                block_table=self.forward_metadata.block_tables,
+                block_size=self.page_size,
+                input_layout="BNSD",
+                scale=layer.scaling,
+                actual_seq_lengths_kv=actual_seq_len_kv,
+                antiquant_mode=0,
+                antiquant_scale=None,
+                sparse_mode=0,
+            )
+            output = torch.zeros_like(q_nope, dtype=q.dtype, device=q.device)
+            softmax_lse = torch.empty(1, dtype=q.dtype, device=q.device)
+
+            torch_npu.npu_fused_infer_attention_score.out(
+                q_nope,
+                c_kv_cache,
+                c_kv_cache,
+                query_rope=q_rope,
+                key_rope=k_rope_cache,
+                num_heads=layer.tp_q_head_num,
+                num_key_value_heads=layer.tp_k_head_num,
+                block_table=self.forward_metadata.block_tables,
+                block_size=self.page_size,
+                input_layout="BNSD",
+                scale=layer.scaling,
+                actual_seq_lengths_kv=actual_seq_len_kv,
+                antiquant_mode=0,
+                antiquant_scale=None,
+                sparse_mode=0,
+                workspace=workspace,
+                out=[output, softmax_lse],
+            )
+            return output.view(-1, layer.tp_q_head_num * self.kv_lora_rank)
+
     def forward_decode(
         self,
         q: torch.Tensor,
@@ -260,106 +425,74 @@ class AscendAttnBackend(AttentionBackend):
         v: torch.Tensor,
         layer: RadixAttention,
         forward_batch: ForwardBatch,
-        save_kv_cache: bool = False,
+        save_kv_cache: bool = True,
         # For multi-head latent attention
         q_rope: Optional[torch.Tensor] = None,
         k_rope: Optional[torch.Tensor] = None,
     ):
+        if self.graph_mode:
+            return self.forward_decode_graph(
+                q,
+                k,
+                v,
+                layer,
+                forward_batch,
+                save_kv_cache,
+                q_rope=q_rope,
+                k_rope=k_rope,
+            )
+
         if not self.use_mla:
             if save_kv_cache:
                 forward_batch.token_to_kv_pool.set_kv_buffer(
                     layer, forward_batch.out_cache_loc, k, v
                 )
             num_tokens = q.shape[0]
-            if self.graph_mode:
-                k_cache = forward_batch.token_to_kv_pool.get_key_buffer(
-                    layer.layer_id
-                ).view(-1, self.page_size, layer.tp_k_head_num * layer.qk_head_dim)
-                v_cache = forward_batch.token_to_kv_pool.get_value_buffer(
-                    layer.layer_id
-                ).view(-1, self.page_size, layer.tp_v_head_num * layer.v_head_dim)
-                query = q.view(-1, 1, layer.tp_q_head_num * layer.qk_head_dim)
-                workspace = (
-                    torch_npu._npu_fused_infer_attention_score_get_max_workspace(
-                        query,
-                        k_cache,
-                        v_cache,
-                        block_table=self.forward_metadata.block_tables,
-                        block_size=self.page_size,
-                        num_heads=layer.tp_q_head_num,
-                        num_key_value_heads=layer.tp_k_head_num,
-                        input_layout="BSH",
-                        scale=layer.scaling,
-                        actual_seq_lengths_kv=self.forward_metadata.seq_lens_cpu_list,
-                    )
-                )
-                attn_output = torch.empty(
-                    (num_tokens, 1, layer.tp_q_head_num * layer.v_head_dim),
-                    dtype=q.dtype,
-                    device=q.device,
-                )
-                softmax_lse = torch.empty(1, dtype=q.dtype, device=q.device)
-                torch_npu.npu_fused_infer_attention_score.out(
-                    query,
-                    k_cache,
-                    v_cache,
-                    block_table=self.forward_metadata.block_tables,
-                    block_size=self.page_size,
+            k_cache = forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id)
+            v_cache = forward_batch.token_to_kv_pool.get_value_buffer(layer.layer_id)
+            if self.use_fia:
+                attn_output, _ = torch.ops.npu.npu_fused_infer_attention_score(
+                    q.view(
+                        forward_batch.batch_size,
+                        -1,
+                        layer.tp_q_head_num,
+                        layer.qk_head_dim,
+                    ),
+                    k_cache.view(
+                        -1, self.page_size, layer.tp_k_head_num * layer.qk_head_dim
+                    ),
+                    v_cache.view(
+                        -1, self.page_size, layer.tp_v_head_num * layer.qk_head_dim
+                    ),
                     num_heads=layer.tp_q_head_num,
                     num_key_value_heads=layer.tp_k_head_num,
-                    input_layout="BSH",
+                    input_layout="BSND",
+                    atten_mask=None,
+                    block_size=self.page_size,
+                    block_table=self.forward_metadata.block_tables,
+                    actual_seq_lengths_kv=self.forward_metadata.seq_lens_cpu_int,
                     scale=layer.scaling,
-                    actual_seq_lengths_kv=self.forward_metadata.seq_lens_cpu_list,
-                    workspace=workspace,
-                    out=[attn_output, softmax_lse],
                 )
             else:
-                k_cache = forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id)
-                v_cache = forward_batch.token_to_kv_pool.get_value_buffer(
-                    layer.layer_id
+                query = q.reshape(-1, layer.tp_q_head_num, layer.qk_head_dim)
+                num_tokens = query.shape[0]
+                attn_output = torch.empty(
+                    (num_tokens, layer.tp_q_head_num, layer.v_head_dim),
+                    dtype=query.dtype,
+                    device=query.device,
                 )
-                if self.use_fia:
-                    attn_output, _ = torch.ops.npu.npu_fused_infer_attention_score(
-                        q.view(
-                            forward_batch.batch_size,
-                            -1,
-                            layer.tp_q_head_num,
-                            layer.qk_head_dim,
-                        ),
-                        k_cache.view(
-                            -1, self.page_size, layer.tp_k_head_num * layer.qk_head_dim
-                        ),
-                        v_cache.view(
-                            -1, self.page_size, layer.tp_v_head_num * layer.qk_head_dim
-                        ),
-                        num_heads=layer.tp_q_head_num,
-                        num_key_value_heads=layer.tp_k_head_num,
-                        input_layout="BSND",
-                        atten_mask=None,
-                        block_size=self.page_size,
-                        block_table=self.forward_metadata.block_tables,
-                        actual_seq_lengths_kv=self.forward_metadata.seq_lens_cpu_int,
-                        scale=layer.scaling,
-                    )
-                else:
-                    query = q.view(-1, layer.tp_q_head_num, layer.qk_head_dim)
-                    attn_output = torch.empty(
-                        (num_tokens, layer.tp_q_head_num, layer.v_head_dim),
-                        dtype=query.dtype,
-                        device=query.device,
-                    )
 
-                    torch_npu._npu_paged_attention(
-                        query=query,
-                        key_cache=k_cache,
-                        value_cache=v_cache,
-                        num_heads=layer.tp_q_head_num,
-                        num_kv_heads=layer.tp_k_head_num,
-                        scale_value=layer.scaling,
-                        block_table=self.forward_metadata.block_tables,
-                        context_lens=self.forward_metadata.seq_lens_cpu_int,
-                        out=attn_output,
-                    )
+                torch_npu._npu_paged_attention(
+                    query=query,
+                    key_cache=k_cache,
+                    value_cache=v_cache,
+                    num_heads=layer.tp_q_head_num,
+                    num_kv_heads=layer.tp_k_head_num,
+                    scale_value=layer.scaling,
+                    block_table=self.forward_metadata.block_tables,
+                    context_lens=self.forward_metadata.seq_lens_cpu_int,
+                    out=attn_output,
+                )
             return attn_output.view(num_tokens, layer.tp_q_head_num * layer.v_head_dim)
         else:
             if save_kv_cache:
@@ -370,9 +503,7 @@ class AscendAttnBackend(AttentionBackend):
             kv_c = forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id)
             k_pe = forward_batch.token_to_kv_pool.get_value_buffer(layer.layer_id)
 
-            if (self.graph_mode or self.use_fia) and (
-                layer.tp_q_head_num // layer.tp_k_head_num
-            ) >= 8:
+            if self.use_fia and (layer.tp_q_head_num // layer.tp_k_head_num) >= 8:
                 """layer.tp_q_head_num // layer.tp_k_head_num < 8 will support in the later version of CANN"""
                 kv_c = kv_c.view(
                     -1, self.page_size, layer.tp_k_head_num * self.kv_lora_rank

sglang/srt/layers/attention/hybrid_attn_backend.py

@@ -5,6 +5,7 @@ import torch
 from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
 from sglang.srt.layers.radix_attention import RadixAttention
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+from sglang.srt.model_executor.model_runner import ModelRunner
 from sglang.srt.speculative.eagle_utils import EagleDraftInput, EagleVerifyInput
 
 
@@ -12,19 +13,27 @@ class HybridAttnBackend(AttentionBackend):
     """Support different backends for prefill and decode."""
 
     def __init__(
-        self, prefill_backend: AttentionBackend, decode_backend: AttentionBackend
+        self,
+        model_runner: ModelRunner,
+        prefill_backend: AttentionBackend,
+        decode_backend: AttentionBackend,
     ):
+        self.model_runner = model_runner
         self.prefill_backend = prefill_backend
         self.decode_backend = decode_backend
 
     def init_forward_metadata(self, forward_batch: ForwardBatch):
-        if forward_batch.forward_mode.is_decode():
+        if forward_batch.forward_mode.is_decode_or_idle():
             self.decode_backend.init_forward_metadata(forward_batch)
         else:
             self.prefill_backend.init_forward_metadata(forward_batch)
 
     def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int):
         self.decode_backend.init_cuda_graph_state(max_bs, max_num_tokens)
+        if self.model_runner.server_args.speculative_algorithm is not None:
+            # When speculative decoding is enabled, we also need to initialize the
+            # prefill backend's cuda graph state to support target_verify.
+            self.prefill_backend.init_cuda_graph_state(max_bs, max_num_tokens)
 
     def init_forward_metadata_capture_cuda_graph(
         self,
@@ -36,15 +45,26 @@ class HybridAttnBackend(AttentionBackend):
         forward_mode: ForwardMode,
         spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]],
     ):
-        self.decode_backend.init_forward_metadata_capture_cuda_graph(
-            bs,
-            num_tokens,
-            req_pool_indices,
-            seq_lens,
-            encoder_lens,
-            forward_mode,
-            spec_info,
-        )
+        if forward_mode.is_decode_or_idle():
+            self.decode_backend.init_forward_metadata_capture_cuda_graph(
+                bs,
+                num_tokens,
+                req_pool_indices,
+                seq_lens,
+                encoder_lens,
+                forward_mode,
+                spec_info,
+            )
+        else:
+            self.prefill_backend.init_forward_metadata_capture_cuda_graph(
+                bs,
+                num_tokens,
+                req_pool_indices,
+                seq_lens,
+                encoder_lens,
+                forward_mode,
+                spec_info,
+            )
 
     def init_forward_metadata_replay_cuda_graph(
         self,
@@ -57,16 +77,28 @@ class HybridAttnBackend(AttentionBackend):
         spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]],
         seq_lens_cpu: Optional[torch.Tensor],
     ):
-        self.decode_backend.init_forward_metadata_replay_cuda_graph(
-            bs,
-            req_pool_indices,
-            seq_lens,
-            seq_lens_sum,
-            encoder_lens,
-            forward_mode,
-            spec_info,
-            seq_lens_cpu,
-        )
+        if forward_mode.is_decode_or_idle():
+            self.decode_backend.init_forward_metadata_replay_cuda_graph(
+                bs,
+                req_pool_indices,
+                seq_lens,
+                seq_lens_sum,
+                encoder_lens,
+                forward_mode,
+                spec_info,
+                seq_lens_cpu,
+            )
+        else:
+            self.prefill_backend.init_forward_metadata_replay_cuda_graph(
+                bs,
+                req_pool_indices,
+                seq_lens,
+                seq_lens_sum,
+                encoder_lens,
+                forward_mode,
+                spec_info,
+                seq_lens_cpu,
+            )
 
     def get_cuda_graph_seq_len_fill_value(self):
         return self.decode_backend.get_cuda_graph_seq_len_fill_value()

sglang/srt/layers/attention/trtllm_mla_backend.py

@@ -51,6 +51,7 @@ class TRTLLMMLADecodeMetadata:
 
     workspace: Optional[torch.Tensor] = None
     block_kv_indices: Optional[torch.Tensor] = None
+    max_seq_len: Optional[int] = None
 
 
 class TRTLLMMLABackend(FlashInferMLAAttnBackend):
@@ -207,8 +208,9 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
             )
 
         # Custom fast-path for decode/idle.
-        max_seqlen_pad = self._calc_padded_blocks(seq_lens.max().item())
-        block_kv_indices = self.decode_cuda_graph_kv_indices[:bs, :max_seqlen_pad]
+        # Capture with full width so future longer sequences are safe during replay
+        max_blocks_per_seq = self._calc_padded_blocks(self.max_context_len)
+        block_kv_indices = self.decode_cuda_graph_kv_indices[:bs, :max_blocks_per_seq]
 
         create_flashmla_kv_indices_triton[(bs,)](
             self.req_to_token,
@@ -217,13 +219,20 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
             None,
             block_kv_indices,
             self.req_to_token.stride(0),
-            max_seqlen_pad,
+            max_blocks_per_seq,
             NUM_PAGE_PER_BLOCK=TRITON_PAD_NUM_PAGE_PER_BLOCK,
             PAGED_SIZE=self.page_size,
         )
 
+        # Record the true maximum sequence length for this capture batch so that
+        # the kernel launch path (which requires an int not a tensor) can reuse
+        # it safely during both capture and replay.
+        max_seq_len_val = int(seq_lens.max().item())
+
         metadata = TRTLLMMLADecodeMetadata(
-            self.decode_cuda_graph_workspace, block_kv_indices
+            self.decode_cuda_graph_workspace,
+            block_kv_indices,
+            max_seq_len_val,
         )
         self.decode_cuda_graph_metadata[bs] = metadata
         self.forward_metadata = metadata
@@ -268,6 +277,13 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
             PAGED_SIZE=self.page_size,
         )
 
+        # Update stored max_seq_len so subsequent kernel calls use the correct value
+        # Prefer CPU tensor to avoid GPU synchronization when available.
+        if seq_lens_cpu is not None:
+            metadata.max_seq_len = int(seq_lens_cpu.max().item())
+        else:
+            metadata.max_seq_len = int(seq_lens.max().item())
+
     def get_cuda_graph_seq_len_fill_value(self) -> int:
         """Get the fill value for sequence lengths in CUDA graph."""
         return 1
@@ -295,8 +311,9 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
             forward_batch.seq_lens.device,
         )
 
+        max_seq_len_val = int(max_seq)
         self.forward_metadata = TRTLLMMLADecodeMetadata(
-            self.workspace_buffer, block_kv_indices
+            self.workspace_buffer, block_kv_indices, max_seq_len_val
         )
         forward_batch.decode_trtllm_mla_metadata = self.forward_metadata
 
@@ -471,14 +488,12 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
             qk_rope_head_dim=self.qk_rope_head_dim,
             block_tables=metadata.block_kv_indices,
             seq_lens=forward_batch.seq_lens.to(torch.int32),
-            max_seq_len=int(metadata.block_kv_indices.shape[1] * self.page_size),
+            max_seq_len=metadata.max_seq_len,
             bmm1_scale=bmm1_scale,
         )
 
-        # Extract value projection part and reshape
-        raw_out_v = raw_out[..., : layer.v_head_dim].contiguous()
-        output = raw_out_v.view(-1, layer.tp_q_head_num * layer.v_head_dim)
-
+        # Reshape output directly without slicing
+        output = raw_out.view(-1, layer.tp_q_head_num * layer.v_head_dim)
         return output
 
 

sglang/srt/layers/layernorm.py

@@ -53,7 +53,7 @@ elif _is_hip:
 
 logger = logging.getLogger(__name__)
 
-if is_npu():
+if _is_npu:
     import torch_npu
 
 
@@ -266,23 +266,48 @@ class GemmaRMSNorm(CustomOp):
         out = gemma_rmsnorm(x, self.weight.data, self.variance_epsilon)
         return out
 
+    def forward_npu(
+        self,
+        x: torch.Tensor,
+        residual: Optional[torch.Tensor] = None,
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        orig_dtype = x.dtype
+        if residual is not None:
+            x = x + residual
+            residual = x
 
-class Gemma3RMSNorm(nn.Module):
+        x = x.float()
+        variance = torch_npu.mean(torch_npu.pow(x, 2), dim=-1, keepdim=True)
+        x = x * torch_npu.rsqrt(variance + self.variance_epsilon)
+        x = x * (1.0 + self.weight.float())
+        x = x.to(orig_dtype)
+        return x if residual is None else (x, residual)
+
+
+class Gemma3RMSNorm(CustomOp):
     def __init__(self, dim: int, eps: float = 1e-6):
         super().__init__()
         self.eps = eps
         self.weight = nn.Parameter(torch.zeros(dim))
+        # Re-dispatch
 
     def _norm(self, x):
         return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
 
-    def forward(self, x):
+    def forward_native(self, x):
         output = self._norm(x.float())
         # Llama does x.to(float16) * w whilst Gemma3 is (x * w).to(float16)
         # See https://github.com/huggingface/transformers/pull/29402
         output = output * (1.0 + self.weight.float())
         return output.type_as(x)
 
+    def forward_cuda(self, x):
+        return self.forward_native(x)
+
+    def forward_npu(self, x):
+        output, _ = torch_npu.npu_gemma_rms_norm(x, self.weight, self.eps)
+        return output
+
     def extra_repr(self):
         return f"{tuple(self.weight.shape)}, eps={self.eps}"