From 79ab57bd7a2a236efd14451fa39607270849d2be Mon Sep 17 00:00:00 2001 From: b8zhong Date: Tue, 16 Dec 2025 23:09:44 -0800 Subject: [PATCH] Revert "direct register custom op for mm_fp4 (#13699)" (#15284) --- .../sglang/srt/layers/attention/trtllm_mha_backend.py | 7 ++++++- .../sglang/srt/layers/quantization/modelopt_quant.py | 11 ++--------- 2 files changed, 8 insertions(+), 10 deletions(-) diff --git a/python/sglang/srt/layers/attention/trtllm_mha_backend.py b/python/sglang/srt/layers/attention/trtllm_mha_backend.py index 8f2447d21..728094dec 100644 --- a/python/sglang/srt/layers/attention/trtllm_mha_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mha_backend.py @@ -529,7 +529,12 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend): forward_batch.extend_prefix_lens_cpu ) or forward_batch.forward_mode.is_draft_extend(include_v2=True): extend_seq_lens = forward_batch.extend_seq_lens - metadata.max_seq_len_q = max(forward_batch.extend_seq_lens_cpu) + # NOTE: in piecewise CUDA graph warmup, extend_seq_lens_cpu is a torch.Tensor; + # Python's max() returns a 0-d tensor, but flashinfer expects an int. + max_q = max(forward_batch.extend_seq_lens_cpu) + metadata.max_seq_len_q = ( + int(max_q.item()) if isinstance(max_q, torch.Tensor) else int(max_q) + ) metadata.cu_seqlens_q = torch.nn.functional.pad( torch.cumsum(extend_seq_lens, dim=0, dtype=torch.int32), (1, 0) ) diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py index c71dd02e2..45ecd113d 100755 --- a/python/sglang/srt/layers/quantization/modelopt_quant.py +++ b/python/sglang/srt/layers/quantization/modelopt_quant.py @@ -47,7 +47,6 @@ from sglang.srt.layers.quantization.utils import ( ) from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.utils.common import ( - direct_register_custom_op, get_bool_env_var, is_cuda, is_sm120_supported, @@ -103,6 +102,7 @@ except ImportError: logger = logging.getLogger(__name__) +@torch.library.custom_op("sglang::fp4_gemm", mutates_args=()) def _sglang_fp4_gemm( input: torch.Tensor, weight: torch.Tensor, @@ -121,6 +121,7 @@ def _sglang_fp4_gemm( return fp4_gemm(input, weight, input_sf, weight_sf, alpha, out_dtype) +@torch.library.register_fake("sglang::fp4_gemm") def _sglang_fp4_gemm_fake( input, weight, @@ -135,14 +136,6 @@ def _sglang_fp4_gemm_fake( return input.new_empty((M, N), dtype=out_dtype) -direct_register_custom_op( - op_name="fp4_gemm", - op_func=_sglang_fp4_gemm, - mutates_args=[], - fake_impl=_sglang_fp4_gemm_fake, -) - - if is_cuda() and (not is_sm120_supported()) and (fp4_quantize is not None): @register_fake_if_exists("sgl_kernel::scaled_fp4_quant")