Enable PDL for various kernels in DSV32/GLM5 (#23965)

Co-authored-by: b8zhong <b8zhong@users.noreply.github.com>

python/sglang/srt/layers/attention/nsa/nsa_indexer.py

@@ -299,6 +299,12 @@ class Indexer(MultiPlatformOp):
         weights = weights.unsqueeze(-1) * q_scale * self.softmax_scale
         return weights
 
+    @torch.compile(dynamic=True)
+    def _apply_q_scale_and_softmax_scale(
+        self, weights: torch.Tensor, q_scale: torch.Tensor
+    ):
+        return weights.unsqueeze(-1) * q_scale * self.softmax_scale
+
     def _get_q_k_bf16(
         self,
         q_lora: torch.Tensor,
@@ -1161,7 +1167,7 @@ class Indexer(MultiPlatformOp):
                     act_quant=act_quant,
                 )
             current_stream.wait_stream(self.alt_stream)
-            weights = weights.unsqueeze(-1) * q_scale * self.softmax_scale
+            weights = self._apply_q_scale_and_softmax_scale(weights, q_scale)
         else:
             query, key = self._get_q_k_bf16(
                 q_lora, x, positions, enable_dual_stream, forward_batch=forward_batch

python/sglang/srt/layers/attention/utils.py

@@ -12,6 +12,8 @@ _is_cuda = is_cuda()
 if _is_cuda:
     from sgl_kernel import concat_mla_absorb_q
 
+from sglang.jit_kernel.utils import is_arch_support_pdl
+
 
 @triton.jit
 def create_flashinfer_kv_indices_triton(
@@ -462,6 +464,7 @@ def mla_quantize_and_rope_for_fp8(
         # Quantization scales (set to 1.0 for no additional scaling)
         quant_scale_q=1.0,
         quant_scale_kv=1.0,
+        enable_pdl=is_arch_support_pdl(),
     )
 
     return q_out, k_nope_out, k_rope_out

python/sglang/srt/mem_cache/utils.py

@@ -20,6 +20,7 @@ import torch
 import triton
 import triton.language as tl
 
+from sglang.jit_kernel.utils import is_arch_support_pdl
 from sglang.srt.environ import envs
 
 
@@ -35,6 +36,7 @@ def set_mla_kv_buffer_kernel(
     nope_dim: tl.constexpr,
     rope_dim: tl.constexpr,
     BLOCK: tl.constexpr,
+    USE_GDC: tl.constexpr = False,
 ):
     pid_loc = tl.program_id(0)
     pid_blk = tl.program_id(1)
@@ -44,6 +46,9 @@ def set_mla_kv_buffer_kernel(
     total_dim = nope_dim + rope_dim
     mask = offs < total_dim
 
+    if USE_GDC:
+        tl.extra.cuda.gdc_wait()
+
     loc = tl.load(loc_ptr + pid_loc).to(tl.int64)
     dst_ptr = kv_buffer_ptr + loc * buffer_stride + offs
 
@@ -82,6 +87,9 @@ def set_mla_kv_buffer_kernel(
 
     tl.store(dst_ptr, src, mask=mask)
 
+    if USE_GDC:
+        tl.extra.cuda.gdc_launch_dependents()
+
 
 def set_mla_kv_buffer_triton(
     kv_buffer: torch.Tensor,
@@ -96,6 +104,8 @@ def set_mla_kv_buffer_triton(
     n_loc = loc.numel()
     grid = (n_loc, triton.cdiv(total_dim, BLOCK))
 
+    pdl_kwargs = {"USE_GDC": True, "launch_pdl": True} if is_arch_support_pdl() else {}
+
     set_mla_kv_buffer_kernel[grid](
         kv_buffer,
         cache_k_nope,
@@ -107,6 +117,7 @@ def set_mla_kv_buffer_triton(
         nope_dim,
         rope_dim,
         BLOCK=BLOCK,
+        **pdl_kwargs,
     )
 
 
@@ -122,6 +133,7 @@ def set_mla_kv_buffer_fp8_quant_kernel(
     nope_dim: tl.constexpr,
     rope_dim: tl.constexpr,
     BLOCK: tl.constexpr,
+    USE_GDC: tl.constexpr = False,
 ):
     """Fuse BF16/FP16->FP8 cast with paged KV write."""
     pid_loc = tl.program_id(0)
@@ -132,6 +144,9 @@ def set_mla_kv_buffer_fp8_quant_kernel(
     total_dim = nope_dim + rope_dim
     mask = offs < total_dim
 
+    if USE_GDC:
+        tl.extra.cuda.gdc_wait()
+
     loc = tl.load(loc_ptr + pid_loc).to(tl.int64)
     dst_ptr = kv_buffer_fp8_ptr + loc * buffer_stride + offs
 
@@ -165,6 +180,9 @@ def set_mla_kv_buffer_fp8_quant_kernel(
     # Destination pointer is FP8-typed view; tl.store performs downcast.
     tl.store(dst_ptr, src, mask=mask)
 
+    if USE_GDC:
+        tl.extra.cuda.gdc_launch_dependents()
+
 
 def set_mla_kv_buffer_triton_fp8_quant(
     kv_buffer: torch.Tensor,
@@ -183,6 +201,8 @@ def set_mla_kv_buffer_triton_fp8_quant(
     n_loc = loc.numel()
     grid = (n_loc, triton.cdiv(total_dim, BLOCK))
 
+    pdl_kwargs = {"USE_GDC": True, "launch_pdl": True} if is_arch_support_pdl() else {}
+
     set_mla_kv_buffer_fp8_quant_kernel[grid](
         kv_buffer_fp8,
         cache_k_nope,
@@ -194,6 +214,7 @@ def set_mla_kv_buffer_triton_fp8_quant(
         nope_dim,
         rope_dim,
         BLOCK=BLOCK,
+        **pdl_kwargs,
     )