Support piecewise CUDA graph with NSA

python/sglang/jit_kernel/hadamard.py

@@ -5,6 +5,7 @@
 import torch
 
 from sglang.jit_kernel.utils import KERNEL_PATH, cache_once, load_jit, make_cpp_args
+from sglang.srt.utils.custom_op import register_custom_op
 
 if TYPE_CHECKING:
     from tvm_ffi.module import Module
@@ -56,6 +57,14 @@ def _hadamard_transform_impl(
     return out.reshape(shapes_og)
 
 
+def _hadamard_transform_fake_impl(
+    x: torch.Tensor,
+    scale: float = 1.0,
+) -> torch.Tensor:
+    return torch.empty_like(x)
+
+
+@register_custom_op(fake_impl=_hadamard_transform_fake_impl)
 def hadamard_transform(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor:
     module = _jit_hadamard_module(x.dtype)
     return _hadamard_transform_impl(x, scale, 8, module.hadamard_transform)

python/sglang/srt/compilation/piecewise_context_manager.py

@@ -71,6 +71,7 @@ def __init__(self):
         self.quant_config = None
         self.moe_layers = None
         self.moe_fusions = None
+        self.dsa_indexers = None
 
     def set_forward_batch(self, forward_batch: ForwardBatch):
         self.forward_batch = forward_batch
@@ -87,6 +88,9 @@ def set_moe_layers(self, layers: List[Any]):
     def set_moe_fusions(self, fusions: List[Any]):
         self.moe_fusions = fusions
 
+    def set_dsa_indexers(self, indexers: List[Any]):
+        self.dsa_indexers = indexers
+
 
 _forward_context: Optional[ForwardContext] = None
 
@@ -104,6 +108,7 @@ def set_forward_context(
     quant_config: Any,
     moe_layers: List[Any],
     moe_fusions: List[Any],
+    dsa_indexers: Optional[List[Any]] = None,
 ):
     global _forward_context
     _forward_context = ForwardContext()
@@ -112,6 +117,8 @@ def set_forward_context(
     _forward_context.set_quant_config(quant_config)
     _forward_context.set_moe_layers(moe_layers)
     _forward_context.set_moe_fusions(moe_fusions)
+    if dsa_indexers is not None:
+        _forward_context.set_dsa_indexers(dsa_indexers)
     try:
         yield
     finally:

python/sglang/srt/configs/model_config.py

@@ -329,7 +329,6 @@ def __init__(
         self.use_ngram_embedding = getattr(self.hf_config, "use_ngram_embedding", False)
         self.is_piecewise_cuda_graph_disabled_model = (
             is_piecewise_cuda_graph_disabled_model(self.hf_config.architectures)
-            or is_deepseek_dsa(self.hf_text_config)
         )
         self.dtype = _get_and_verify_dtype(self.hf_text_config, dtype)
 
@@ -1556,11 +1555,9 @@ def is_generation_model(model_architectures: List[str], is_embedding: bool = Fal
 ]
 
 piecewise_cuda_graph_disabled_model_archs = [
-    "DeepseekV32ForCausalLM",
     "DeepseekV4ForCausalLM",
     "DeepseekV4ForCausalLMNextN",
     "Qwen3NextForCausalLM",
-    "GlmMoeDsaForCausalLM",
     "BailingMoeV2_5ForCausalLM",
     "LLaDAModelLM",
 ]