sglang 0.5.2rc2__py3-none-any.whl → 0.5.3__py3-none-any.whl

sglang/srt/layers/attention/trtllm_mla_backend.py

@@ -20,8 +20,9 @@ from sglang.srt.layers.attention.utils import (
     create_flashmla_kv_indices_triton,
 )
 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
@@ -29,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
@@ -45,11 +51,19 @@ TRTLLM_BLOCK_CONSTRAINT = 128
 global_zero_init_workspace_buffer = None
 
 
+@dataclass
+class TRTLLMMLAPrefillMetadata:
+    """Metadata for TRTLLM MLA prefill operations."""
+
+    max_seq_len: int
+    cum_seq_lens: torch.Tensor
+    seq_lens: torch.Tensor
+
+
 @dataclass
 class TRTLLMMLADecodeMetadata:
     """Metadata for TRTLLM MLA decode operations."""
 
-    workspace: Optional[torch.Tensor] = None
     block_kv_indices: Optional[torch.Tensor] = None
     max_seq_len: Optional[int] = None
 
@@ -64,7 +78,12 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
         kv_indptr_buf: Optional[torch.Tensor] = None,
         q_indptr_decode_buf: Optional[torch.Tensor] = None,
     ):
-        super().__init__(model_runner, skip_prefill, kv_indptr_buf, q_indptr_decode_buf)
+        super().__init__(
+            model_runner,
+            skip_prefill,
+            kv_indptr_buf,
+            q_indptr_decode_buf,
+        )
 
         config = model_runner.model_config
 
@@ -101,7 +120,14 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
         # CUDA graph state
         self.decode_cuda_graph_metadata = {}
         self.decode_cuda_graph_kv_indices = None
-        self.forward_metadata: Union[TRTLLMMLADecodeMetadata, None] = None
+        self.forward_prefill_metadata: Optional[TRTLLMMLAPrefillMetadata] = None
+        self.forward_decode_metadata: Union[TRTLLMMLADecodeMetadata, None] = None
+
+        self.disable_chunked_prefix_cache = global_server_args_dict[
+            "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:
         """
@@ -177,9 +203,6 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
         self.decode_cuda_graph_kv_indices = torch.full(
             (max_bs, max_blocks_per_seq), -1, dtype=torch.int32, device=self.device
         )
-        self.decode_cuda_graph_workspace = torch.empty(
-            self.workspace_size, dtype=torch.int8, device=self.device
-        )
 
         super().init_cuda_graph_state(max_bs, max_num_tokens, kv_indices_buf)
 
@@ -191,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,
@@ -207,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)
@@ -230,12 +256,11 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
         max_seq_len_val = int(seq_lens.max().item())
 
         metadata = TRTLLMMLADecodeMetadata(
-            self.decode_cuda_graph_workspace,
             block_kv_indices,
             max_seq_len_val,
         )
         self.decode_cuda_graph_metadata[bs] = metadata
-        self.forward_metadata = metadata
+        self.forward_decode_metadata = metadata
 
     def init_forward_metadata_replay_cuda_graph(
         self,
@@ -245,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,
@@ -262,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.
@@ -291,31 +320,64 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
     def init_forward_metadata(self, forward_batch: ForwardBatch):
         """Initialize the metadata for a forward pass."""
         # Delegate to parent for non-decode modes.
-        if not forward_batch.forward_mode.is_decode_or_idle():
-            return super().init_forward_metadata(forward_batch)
-
-        bs = forward_batch.batch_size
+        if (
+            forward_batch.forward_mode.is_extend()
+            and not forward_batch.forward_mode.is_target_verify()
+            and not forward_batch.forward_mode.is_draft_extend()
+        ):
+            if self.disable_chunked_prefix_cache:
+                super().init_forward_metadata(forward_batch)
+
+            seq_lens = forward_batch.seq_lens - forward_batch.extend_prefix_lens
+            cum_seq_lens_q = torch.cat(
+                (
+                    torch.tensor([0], device=forward_batch.seq_lens.device),
+                    torch.cumsum(seq_lens, dim=0),
+                )
+            ).int()
+            max_seq_len = max(forward_batch.extend_seq_lens_cpu)
+            self.forward_prefill_metadata = TRTLLMMLAPrefillMetadata(
+                max_seq_len,
+                cum_seq_lens_q,
+                seq_lens,
+            )
+        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.
+            if getattr(forward_batch, "seq_lens_cpu", None) is not None:
+                max_seq = forward_batch.seq_lens_cpu.max().item()
+            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,
+                seq_lens,
+                seq_lens.device,
+            )
 
-        # Get maximum sequence length.
-        if getattr(forward_batch, "seq_lens_cpu", None) is not None:
-            max_seq = forward_batch.seq_lens_cpu.max().item()
+            max_seq_len_val = int(max_seq)
+            self.forward_decode_metadata = TRTLLMMLADecodeMetadata(
+                block_kv_indices, max_seq_len_val
+            )
+            forward_batch.decode_trtllm_mla_metadata = self.forward_decode_metadata
         else:
-            max_seq = forward_batch.seq_lens.max().item()
-
-        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,
-        )
+            return super().init_forward_metadata(forward_batch)
 
-        max_seq_len_val = int(max_seq)
-        self.forward_metadata = TRTLLMMLADecodeMetadata(
-            self.workspace_buffer, block_kv_indices, max_seq_len_val
-        )
-        forward_batch.decode_trtllm_mla_metadata = self.forward_metadata
+    def init_mha_chunk_metadata(self, forward_batch: ForwardBatch):
+        super().init_mha_chunk_metadata(forward_batch, disable_flashinfer_ragged=True)
 
     def quantize_and_rope_for_fp8(
         self,
@@ -443,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)
@@ -459,7 +521,7 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
         # Get metadata
         metadata = (
             getattr(forward_batch, "decode_trtllm_mla_metadata", None)
-            or self.forward_metadata
+            or self.forward_decode_metadata
         )
 
         # Scale computation for TRTLLM MLA kernel BMM1 operation:
@@ -482,7 +544,7 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
         raw_out = flashinfer.decode.trtllm_batch_decode_with_kv_cache_mla(
             query=query,
             kv_cache=kv_cache,
-            workspace_buffer=metadata.workspace,
+            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,
@@ -496,6 +558,174 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
         output = raw_out.view(-1, layer.tp_q_head_num * layer.v_head_dim)
         return output
 
+    def forward_extend(
+        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,
+        cos_sin_cache: Optional[torch.Tensor] = None,
+        is_neox: Optional[bool] = False,
+    ) -> torch.Tensor:
+        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
+            )
+
+        # 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
+            )
+
+        # 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)
+
+        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=forward_batch.prefix_chunk_seq_lens[chunk_idx],
+                max_q_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,
+                batch_size=forward_batch.batch_size,
+                window_left=-1,
+                cum_seq_lens_q=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=False,
+                return_lse=True,
+                out=torch.zeros(*output_shape, dtype=q.dtype, device=q.device),
+            )
+
+        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):
     """Multi-step draft backend for TRT-LLM MLA used by EAGLE."""
@@ -512,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/attention/wave_ops/decode_attention.py

@@ -64,8 +64,7 @@ def get_wave_kernel(
         subs=hyperparams_0,
         canonicalize=True,
         run_bench=False,
-        use_buffer_load_ops=True,
-        use_buffer_store_ops=True,
+        use_buffer_ops=True,
         waves_per_eu=2,
         dynamic_symbols=dynamic_symbols_0,
         wave_runtime=True,
@@ -77,8 +76,7 @@ def get_wave_kernel(
         subs=hyperparams_1,
         canonicalize=True,
         run_bench=False,
-        use_buffer_load_ops=False,
-        use_buffer_store_ops=False,
+        use_buffer_ops=False,
         waves_per_eu=4,
         dynamic_symbols=dynamic_symbols_1,
         wave_runtime=True,

sglang/srt/layers/attention/wave_ops/extend_attention.py

@@ -67,11 +67,9 @@ def get_wave_kernel(
         schedule=SchedulingType.NONE,
         use_scheduling_barriers=False,
         dynamic_symbols=dynamic_symbols,
-        use_buffer_load_ops=True,
-        use_buffer_store_ops=True,
+        use_buffer_ops=True,
         waves_per_eu=2,
         denorm_fp_math_f32="preserve-sign",
-        gpu_native_math_precision=True,
         wave_runtime=True,
     )
     options = set_default_run_config(options)

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