sglang 0.5.3rc0__py3-none-any.whl → 0.5.4__py3-none-any.whl

sglang/srt/layers/attention/ascend_backend.py

@@ -5,21 +5,21 @@ from typing import TYPE_CHECKING, List, Optional
 
 import torch
 import torch_npu
-from torch.nn.functional import scaled_dot_product_attention
 
 from sglang.srt.configs.model_config import AttentionArch
 from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
+from sglang.srt.layers.attention.npu_ops.mla_preprocess import is_mla_preprocess_enabled
 from sglang.srt.layers.attention.torch_native_backend import TorchNativeAttnBackend
 from sglang.srt.layers.dp_attention import get_attention_tp_size
 from sglang.srt.layers.radix_attention import AttentionType
-from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+from sglang.srt.speculative.spec_info import SpecInput
 from sglang.srt.utils import get_bool_env_var
 
 if TYPE_CHECKING:
     from sglang.srt.layers.radix_attention import RadixAttention
     from sglang.srt.model_executor.model_runner import ModelRunner
 
-import os
 
 import numpy as np
 
@@ -35,6 +35,8 @@ class ForwardMetadata:
     seq_lens_cpu_int: Optional[torch.Tensor] = None
     seq_lens_cpu_list: Optional[List[int]] = None
     seq_lens_list_cumsum: Optional[List[int]] = None
+    seq_lens: Optional[torch.Tensor] = None
+    actual_seq_lengths_q: Optional[torch.Tensor] = None
 
 
 class AscendAttnBackend(AttentionBackend):
@@ -66,6 +68,9 @@ 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.q_head_dim = (
+                self.qk_rope_head_dim + model_runner.model_config.qk_nope_head_dim
+            )
         self.native_attn = TorchNativeAttnBackend(model_runner)
         self.graph_metadata = {}
         self.max_context_len = model_runner.model_config.context_len
@@ -101,10 +106,6 @@ class AscendAttnBackend(AttentionBackend):
         self.forward_metadata.seq_lens_cpu_int = forward_batch.seq_lens_cpu.int()
 
         seq_lens_list_cumsum = np.cumsum(forward_batch.extend_seq_lens_cpu)
-        if forward_batch.is_extend_in_batch:
-            seq_lens_list_cumsum[-1] = (
-                (seq_lens_list_cumsum[-1] - 1) // tp_size + 1
-            ) * tp_size
         self.forward_metadata.seq_lens_list_cumsum = seq_lens_list_cumsum
 
         self.graph_mode = False
@@ -126,12 +127,16 @@ class AscendAttnBackend(AttentionBackend):
         seq_lens: torch.Tensor,
         encoder_lens: Optional[torch.Tensor],
         forward_mode: ForwardMode,
-        spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]],
+        spec_info: Optional[SpecInput],
     ):
         metadata = ForwardMetadata()
 
         metadata.block_tables = self.graph_metadata["block_tables"][:bs, :]
         metadata.seq_lens_cpu_list = seq_lens.cpu().int().tolist()
+        metadata.seq_lens = seq_lens
+        metadata.actual_seq_lengths_q = torch.tensor(
+            [1 + i * 1 for i in range(bs)], dtype=torch.int32, device=seq_lens.device
+        )
 
         self.graph_metadata[bs] = metadata
         self.forward_metadata = metadata
@@ -146,7 +151,7 @@ class AscendAttnBackend(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],
     ):
         metadata = self.graph_metadata[bs]
@@ -160,6 +165,8 @@ class AscendAttnBackend(AttentionBackend):
         metadata.block_tables[:bs, max_seq_pages:].fill_(0)
         metadata.block_tables[bs:, :].fill_(0)
 
+        metadata.seq_lens[:bs].copy_(seq_lens[:bs])
+
         self.forward_metadata = metadata
 
         self.graph_mode = True
@@ -167,6 +174,64 @@ class AscendAttnBackend(AttentionBackend):
     def get_cuda_graph_seq_len_fill_value(self):
         return 0
 
+    def forward_sparse(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        layer: RadixAttention,
+        forward_batch: ForwardBatch,
+        save_kv_cache: bool = True,
+        # For multi_head latent attention
+        q_rope: Optional[torch.Tensor] = None,
+        k_rope: Optional[torch.Tensor] = None,
+        topk_indices: torch.Tensor = None,
+    ):
+
+        is_prefill = forward_batch.forward_mode.is_extend()
+
+        if save_kv_cache:
+            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
+            )
+        q_nope, q_pe = q, q_rope
+        k_nope, k_pe = forward_batch.token_to_kv_pool.get_kv_buffer(layer.layer_id)
+        block_table = self.forward_metadata.block_tables
+        if is_prefill:
+            actual_seq_qlen = torch.cumsum(forward_batch.seq_lens, dim=0)
+        else:
+            if self.forward_metadata.actual_seq_lengths_q is None:
+                actual_seq_qlen = (
+                    torch.arange(1, q.shape[0] + 1).to(q.device).to(torch.int32)
+                )
+            else:
+                actual_seq_qlen = self.forward_metadata.actual_seq_lengths_q
+        if self.forward_metadata.seq_lens_cpu_int is None:
+            actual_seq_lengths_kv = self.forward_metadata.seq_lens
+        else:
+            actual_seq_lengths_kv = self.forward_metadata.seq_lens_cpu_int
+
+        attn_out = torch.ops.custom.npu_sparse_flash_attention(
+            query=q_nope,
+            key=k_nope,
+            value=k_nope,
+            query_rope=q_pe,
+            key_rope=k_pe,
+            sparse_indices=topk_indices,
+            scale_value=layer.scaling,
+            actual_seq_lengths_query=actual_seq_qlen.to(torch.int32),
+            actual_seq_lengths_kv=actual_seq_lengths_kv.to(q.device),
+            block_table=block_table,
+            sparse_block_size=1,
+            layout_query="TND",
+            layout_kv="PA_BSND",
+            sparse_mode=3,
+        )
+
+        return attn_out
+
     def forward_extend(
         self,
         q,
@@ -175,7 +240,23 @@ class AscendAttnBackend(AttentionBackend):
         layer: RadixAttention,
         forward_batch: ForwardBatch,
         save_kv_cache: bool = True,
+        # For multi_head latent attention
+        q_rope: Optional[torch.Tensor] = None,
+        k_rope: Optional[torch.Tensor] = None,
+        topk_indices: Optional[torch.Tensor] = None,
     ):
+        if topk_indices is not None:
+            return self.forward_sparse(
+                q,
+                k,
+                v,
+                layer,
+                forward_batch,
+                save_kv_cache,
+                q_rope,
+                k_rope,
+                topk_indices,
+            )
         if not self.use_mla:
             if save_kv_cache:
                 forward_batch.token_to_kv_pool.set_kv_buffer(
@@ -274,6 +355,11 @@ class AscendAttnBackend(AttentionBackend):
             assert (
                 layer.qk_head_dim != layer.v_head_dim
             ), "FIA only supports qk_head_dim != v_head_dim"
+            num_token_padding = q.shape[0]
+            q, k, v = [
+                data[: forward_batch.num_token_non_padded_cpu] for data in [q, k, v]
+            ]
+
             q_nope, q_rope = q.split([layer.v_head_dim, self.qk_rope_head_dim], dim=-1)
             k_nope, k_rope = k.split([layer.v_head_dim, self.qk_rope_head_dim], dim=-1)
 
@@ -293,6 +379,18 @@ class AscendAttnBackend(AttentionBackend):
                 next_tokens=0,
             )
 
+            attn_output = attn_output.reshape(-1, layer.tp_q_head_num, layer.v_head_dim)
+            if num_token_padding != forward_batch.num_token_non_padded_cpu:
+                attn_output = torch.cat(
+                    [
+                        attn_output,
+                        attn_output.new_zeros(
+                            num_token_padding - attn_output.shape[0],
+                            *attn_output.shape[1:],
+                        ),
+                    ],
+                    dim=0,
+                )
         return attn_output
 
     def forward_decode_graph(
@@ -401,7 +499,7 @@ class AscendAttnBackend(AttentionBackend):
                 antiquant_scale=None,
                 sparse_mode=0,
             )
-            output = torch.zeros_like(q_nope, dtype=q.dtype, device=q.device)
+            output = torch.empty_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(
@@ -436,7 +534,24 @@ class AscendAttnBackend(AttentionBackend):
         # For multi-head latent attention
         q_rope: Optional[torch.Tensor] = None,
         k_rope: Optional[torch.Tensor] = None,
+        topk_indices: Optional[torch.Tensor] = None,
     ):
+        if is_mla_preprocess_enabled():
+            # MLAPO does saving kv_cache
+            save_kv_cache = False
+        if topk_indices is not None:
+            return self.forward_sparse(
+                q,
+                k,
+                v,
+                layer,
+                forward_batch,
+                save_kv_cache,
+                q_rope,
+                k_rope,
+                topk_indices,
+            )
+
         if self.graph_mode:
             return self.forward_decode_graph(
                 q,

sglang/srt/layers/attention/attention_registry.py

@@ -0,0 +1,226 @@
+import logging
+from typing import TYPE_CHECKING
+
+logger = logging.getLogger(__name__)
+
+
+if TYPE_CHECKING:
+    # evade circular imports
+    from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
+    from sglang.srt.model_executor.model_runner import ModelRunner
+
+ATTENTION_BACKENDS = {}
+
+
+def register_attention_backend(name):
+    def decorator(fn):
+        ATTENTION_BACKENDS[name] = fn
+        return fn
+
+    return decorator
+
+
+@register_attention_backend("flashinfer")
+def create_flashinfer_backend(runner):
+    import torch
+
+    if not runner.use_mla_backend:
+        from sglang.srt.layers.attention.flashinfer_backend import FlashInferAttnBackend
+
+        # Init streams
+        if runner.server_args.speculative_algorithm == "EAGLE":
+            if (
+                not hasattr(runner, "plan_stream_for_flashinfer")
+                or not runner.plan_stream_for_flashinfer
+            ):
+                runner.plan_stream_for_flashinfer = torch.cuda.Stream()
+        return FlashInferAttnBackend(
+            runner, init_new_workspace=runner.init_new_workspace
+        )
+    else:
+        from sglang.srt.layers.attention.flashinfer_mla_backend import (
+            FlashInferMLAAttnBackend,
+        )
+
+        return FlashInferMLAAttnBackend(runner)
+
+
+@register_attention_backend("trtllm_mla")
+def create_trtllm_mla_backend(runner):
+    if not runner.use_mla_backend:
+        raise ValueError("trtllm_mla backend can only be used with MLA models.")
+    from sglang.srt.layers.attention.trtllm_mla_backend import TRTLLMMLABackend
+
+    return TRTLLMMLABackend(runner)
+
+
+@register_attention_backend("aiter")
+def create_aiter_backend(runner):
+    from sglang.srt.layers.attention.aiter_backend import AiterAttnBackend
+
+    return AiterAttnBackend(runner)
+
+
+@register_attention_backend("wave")
+def create_wave_backend(runner):
+    from sglang.srt.layers.attention.wave_backend import WaveAttnBackend
+
+    return WaveAttnBackend(runner)
+
+
+@register_attention_backend("ascend")
+def create_ascend_backend(runner):
+    from sglang.srt.layers.attention.ascend_backend import AscendAttnBackend
+
+    return AscendAttnBackend(runner)
+
+
+@register_attention_backend("nsa")
+def create_nsa_backend(runner):
+    from sglang.srt.layers.attention.nsa_backend import NativeSparseAttnBackend
+
+    return NativeSparseAttnBackend(runner)
+
+
+@register_attention_backend("triton")
+def create_triton_backend(runner):
+    assert not runner.model_config.is_encoder_decoder, (
+        "Cross attention is not supported in the triton attention backend. "
+        "Please use `--attention-backend flashinfer`."
+    )
+    if runner.server_args.enable_double_sparsity:
+        from sglang.srt.layers.attention.double_sparsity_backend import (
+            DoubleSparseAttnBackend,
+        )
+
+        return DoubleSparseAttnBackend(runner)
+    else:
+        from sglang.srt.layers.attention.triton_backend import TritonAttnBackend
+
+        return TritonAttnBackend(runner)
+
+
+@register_attention_backend("torch_native")
+def create_torch_native_backend(runner):
+    from sglang.srt.layers.attention.torch_native_backend import TorchNativeAttnBackend
+
+    return TorchNativeAttnBackend(runner)
+
+
+@register_attention_backend("flex_attention")
+def create_flex_attention_backend(runner):
+    from sglang.srt.layers.attention.torch_flex_backend import TorchFlexAttnBackend
+
+    return TorchFlexAttnBackend(runner)
+
+
+@register_attention_backend("flashmla")
+def create_flashmla_backend(runner):
+    from sglang.srt.layers.attention.flashmla_backend import FlashMLABackend
+
+    return FlashMLABackend(runner)
+
+
+@register_attention_backend("fa3")
+def create_flashattention_v3_backend(runner):
+    import torch
+
+    assert (
+        torch.cuda.get_device_capability()[0] == 8 and not runner.use_mla_backend
+    ) or torch.cuda.get_device_capability()[0] == 9, (
+        "FlashAttention v3 Backend requires SM>=80 and SM<=90. "
+        "Please use `--attention-backend flashinfer`."
+    )
+    from sglang.srt.layers.attention.flashattention_backend import FlashAttentionBackend
+
+    return FlashAttentionBackend(runner)
+
+
+@register_attention_backend("fa4")
+def create_flashattention_v4_backend(runner):
+    from sglang.srt.layers.attention.flashattention_backend import FlashAttentionBackend
+
+    return FlashAttentionBackend(runner, fa_impl_ver=4)
+
+
+@register_attention_backend("cutlass_mla")
+def create_cutlass_mla_backend(runner):
+    from sglang.srt.layers.attention.cutlass_mla_backend import CutlassMLABackend
+
+    return CutlassMLABackend(runner)
+
+
+@register_attention_backend("trtllm_mha")
+def create_trtllm_mha_backend(runner):
+    if runner.use_mla_backend:
+        raise ValueError("trtllm_mha backend can only be used with non-MLA models.")
+    from sglang.srt.layers.attention.trtllm_mha_backend import TRTLLMHAAttnBackend
+
+    return TRTLLMHAAttnBackend(runner)
+
+
+@register_attention_backend("intel_amx")
+def create_intel_amx_backend(runner):
+    from sglang.srt.layers.attention.intel_amx_backend import IntelAMXAttnBackend
+
+    return IntelAMXAttnBackend(runner)
+
+
+@register_attention_backend("dual_chunk_flash_attn")
+def create_dual_chunk_flash_attn_backend(runner):
+    from sglang.srt.layers.attention.dual_chunk_flashattention_backend import (
+        DualChunkFlashAttentionBackend,
+    )
+
+    return DualChunkFlashAttentionBackend(runner)
+
+
+def attn_backend_wrapper(runner: "ModelRunner", full_attn_backend: "AttentionBackend"):
+    """
+    Wrapper for special models like hybrid GDN, so we don't
+    need to change the code of the original attention backend.
+    """
+    assert not (
+        runner.hybrid_gdn_config is not None and runner.use_mla_backend
+    ), "hybrid_gdn can only be used with non-MLA models."
+
+    if cfg := runner.mambaish_config:
+        from sglang.srt.layers.attention.fla.utils import check_environments
+        from sglang.srt.layers.attention.hybrid_linear_attn_backend import (
+            GDNAttnBackend,
+            HybridLinearAttnBackend,
+            Mamba2AttnBackend,
+        )
+        from sglang.srt.utils import is_blackwell, is_npu
+
+        check_environments()
+        if runner.hybrid_gdn_config is not None:
+            if is_blackwell():
+                assert (
+                    runner.server_args.attention_backend == "triton"
+                ), "triton backend is the only supported backend on Blackwell GPUs for hybrid GDN models, use --attention-backend triton to specify the backend."
+            if is_npu():
+                assert (
+                    runner.server_args.attention_backend == "ascend"
+                ), "ascend backend is the only supported backend on NPU for hybrid GDN models, use --attention-backend ascend to specify the backend."
+            logger.info(f"Using hybrid linear attention backend for hybrid GDN models.")
+            linear_attn_backend = GDNAttnBackend(runner)
+        elif runner.mamba2_config is not None:
+            linear_attn_backend = Mamba2AttnBackend(runner)
+        else:
+            raise ValueError(
+                "Expected hybrid GDN or NemotronH models, but got unknown model."
+            )
+        full_attn_layers = cfg.full_attention_layer_ids
+        return HybridLinearAttnBackend(
+            full_attn_backend, linear_attn_backend, full_attn_layers
+        )
+
+    return full_attn_backend
+
+
+@register_attention_backend("intel_xpu")
+def create_intel_xpu_backend(runner):
+    from sglang.srt.layers.attention.xpu_backend import XPUAttentionBackend
+
+    return XPUAttentionBackend(runner)

sglang/srt/layers/attention/base_attn_backend.py

@@ -1,14 +1,15 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Optional, Union
+from typing import TYPE_CHECKING, Optional
 
 import torch
 
 if TYPE_CHECKING:
+    from sglang.srt.layers.attention.nsa.nsa_indexer import BaseIndexerMetadata
     from sglang.srt.layers.radix_attention import RadixAttention
     from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
-    from sglang.srt.speculative.eagle_utils import EagleDraftInput, EagleVerifyInput
+    from sglang.srt.speculative.spec_info import SpecInput
 
 
 class AttentionBackend(ABC):
@@ -31,7 +32,7 @@ class AttentionBackend(ABC):
         seq_lens: torch.Tensor,
         encoder_lens: Optional[torch.Tensor],
         forward_mode: ForwardMode,
-        spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]],
+        spec_info: Optional[SpecInput],
     ):
         """Init the metadata for a forward pass for capturing a cuda graph."""
         raise NotImplementedError()
@@ -44,7 +45,7 @@ class AttentionBackend(ABC):
         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],
     ):
         """Init the metadata for a forward pass for replaying a cuda graph."""
@@ -54,6 +55,25 @@ class AttentionBackend(ABC):
         """Get the fill value for padded seq lens. Typically, it is 0 or 1."""
         raise NotImplementedError()
 
+    def get_verify_buffers_to_fill_after_draft(self):
+        """
+        Return buffers of verify attention kernels that needs to be filled after draft.
+
+        Typically, these are tree mask and position buffers.
+        """
+        return [None, None]
+
+    def update_verify_buffers_to_fill_after_draft(
+        self, spec_info: SpecInput, cuda_graph_bs: Optional[int]
+    ):
+        """
+        Update the buffers returned by get_verify_fill_after_draft_buffers if needed.
+
+        Here, we need to redo the computation of all metadata of the attention backend
+        that depends on tree mask and position buffers.
+        """
+        raise NotImplementedError()
+
     def forward(
         self,
         q: torch.Tensor,
@@ -115,3 +135,11 @@ class AttentionBackend(ABC):
     def support_triton(self):
         """Check if the current backend supports triton."""
         return True
+
+    def get_indexer_metadata(
+        self,
+        layer_id: int,
+        forward_batch: ForwardBatch,
+    ) -> Optional[BaseIndexerMetadata]:
+        """Get the indexer metadata. None means don't support indexer."""
+        return None

sglang/srt/layers/attention/cutlass_mla_backend.py

@@ -20,7 +20,7 @@ from sglang.srt.utils import is_cuda
 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:
@@ -151,7 +151,7 @@ class CutlassMLABackend(FlashInferMLAAttnBackend):
         seq_lens: torch.Tensor,
         encoder_lens: Optional[torch.Tensor],
         forward_mode: ForwardMode,
-        spec_info: Optional[SpecInfo],
+        spec_info: Optional[SpecInput],
     ):
         if forward_mode.is_decode_or_idle():
             if spec_info is None:
@@ -190,7 +190,7 @@ class CutlassMLABackend(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],
     ):
 

sglang/srt/layers/attention/double_sparsity_backend.py

@@ -5,8 +5,8 @@ from typing import TYPE_CHECKING
 import torch
 
 from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
-from sglang.srt.managers.schedule_batch import global_server_args_dict
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+from sglang.srt.server_args import get_global_server_args
 
 if TYPE_CHECKING:
     from sglang.srt.layers.radix_attention import RadixAttention
@@ -42,7 +42,7 @@ class DoubleSparseAttnBackend(AttentionBackend):
         # TODO: Change the hard-coded block_seq_num
         self.BLOCK_SEQ = 128
 
-        if global_server_args_dict.get("triton_attention_reduce_in_fp32", False):
+        if get_global_server_args().triton_attention_reduce_in_fp32:
             self.reduce_dtype = torch.float32
         else:
             self.reduce_dtype = torch.float16

sglang/srt/layers/attention/dual_chunk_flashattention_backend.py

@@ -1537,7 +1537,7 @@ class DualChunkFlashAttentionBackend(AttentionBackend):
                     query_inter,
                     key_cache,
                     value_cache,
-                    block_table[:, : decode_meta.max_seq_len_inter],
+                    block_table,
                     decode_meta.seq_lens_inter,
                     softmax_scale,
                     causal=False,

sglang/srt/layers/attention/fla/chunk.py

@@ -2,7 +2,6 @@
 # -*- coding: utf-8 -*-
 # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
 
-import warnings
 from typing import Optional
 
 import torch

sglang/srt/layers/attention/fla/chunk_o.py

@@ -2,7 +2,7 @@
 # -*- coding: utf-8 -*-
 # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
 
-from typing import Optional, Tuple
+from typing import Optional
 
 import torch
 import triton

sglang/srt/layers/attention/fla/chunk_scaled_dot_kkt.py

@@ -74,8 +74,7 @@ def chunk_scaled_dot_kkt_fwd_kernel(
             (1, 0),
         )
         b_k = tl.load(p_k, boundary_check=(0, 1))
-        b_kb = b_k * b_beta[:, None]
-        b_A += tl.dot(b_kb.to(b_k.dtype), tl.trans(b_k))
+        b_A += tl.dot(b_k, tl.trans(b_k))
 
     if USE_G:
         p_g = tl.make_block_ptr(
@@ -85,6 +84,7 @@ def chunk_scaled_dot_kkt_fwd_kernel(
         b_g_diff = b_g[:, None] - b_g[None, :]
         b_A = b_A * safe_exp(b_g_diff)
 
+    b_A *= b_beta[:, None]
     b_A = tl.where(o_t[:, None] > o_t[None, :], b_A, 0)
     p_A = tl.make_block_ptr(
         A + (bos * H + i_h) * BT, (T, BT), (BT * H, 1), (i_t * BT, 0), (BT, BT), (1, 0)

sglang/srt/layers/attention/fla/fused_recurrent.py

@@ -86,8 +86,8 @@ def fused_recurrent_gated_delta_rule_fwd_kernel(
         b_g = tl.load(p_g).to(tl.float32)
 
         if USE_QK_L2NORM_IN_KERNEL:
-            b_q = b_q / (tl.sqrt(tl.sum(b_q * b_q)) + 1e-6)
-            b_k = b_k / (tl.sqrt(tl.sum(b_k * b_k)) + 1e-6)
+            b_q = b_q / (tl.sqrt(tl.sum(b_q * b_q) + 1e-6))
+            b_k = b_k / (tl.sqrt(tl.sum(b_k * b_k) + 1e-6))
         b_q = b_q * scale
         # [BK, BV]
         b_h *= exp(b_g)
@@ -411,8 +411,8 @@ def fused_recurrent_gated_delta_rule_update_fwd_kernel(
         b_g = tl.load(p_g).to(tl.float32)
 
         if USE_QK_L2NORM_IN_KERNEL:
-            b_q = b_q / (tl.sqrt(tl.sum(b_q * b_q)) + 1e-6)
-            b_k = b_k / (tl.sqrt(tl.sum(b_k * b_k)) + 1e-6)
+            b_q = b_q / (tl.sqrt(tl.sum(b_q * b_q) + 1e-6))
+            b_k = b_k / (tl.sqrt(tl.sum(b_k * b_k) + 1e-6))
         b_q = b_q * scale
         # [BK, BV]
         b_h *= exp(b_g)

sglang/srt/layers/attention/fla/fused_sigmoid_gating_recurrent.py

@@ -119,8 +119,8 @@ def fused_sigmoid_gating_delta_rule_update_kernel(
 
         # Apply L2 normalization if enabled
         if USE_QK_L2NORM_IN_KERNEL:
-            b_q = b_q / (tl.sqrt(tl.sum(b_q * b_q)) + 1e-6)
-            b_k = b_k / (tl.sqrt(tl.sum(b_k * b_k)) + 1e-6)
+            b_q = b_q / (tl.sqrt(tl.sum(b_q * b_q) + 1e-6))
+            b_k = b_k / (tl.sqrt(tl.sum(b_k * b_k) + 1e-6))
 
         b_q = b_q * scale
 

sglang/srt/layers/attention/fla/index.py

@@ -3,9 +3,7 @@
 # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
 
 import torch
-import torch.nn.functional as F
 import triton
-import triton.language as tl
 
 from sglang.srt.layers.attention.fla.utils import tensor_cache