minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Revert "direct register custom op for mm_fp4 (#13699)" (#15284)

Browse Source
This commit is contained in:
b8zhong
2025-12-16 23:09:44 -08:00
committed by GitHub
parent 4b8901ac0f
commit 79ab57bd7a
2 changed files with 8 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/trtllm_mha_backend.py
+6 -1
View File
@@ -529,7 +529,12 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
forward_batch.extend_prefix_lens_cpu forward_batch.extend_prefix_lens_cpu
) or forward_batch.forward_mode.is_draft_extend(include_v2=True): ) or forward_batch.forward_mode.is_draft_extend(include_v2=True):
extend_seq_lens = forward_batch.extend_seq_lens 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( metadata.cu_seqlens_q = torch.nn.functional.pad(
torch.cumsum(extend_seq_lens, dim=0, dtype=torch.int32), (1, 0) torch.cumsum(extend_seq_lens, dim=0, dtype=torch.int32), (1, 0)
) )
python/sglang/srt/layers/quantization/modelopt_quant.py
+2 -9
View File
@@ -47,7 +47,6 @@ from sglang.srt.layers.quantization.utils import (
) )
from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.layers.radix_attention import RadixAttention
from sglang.srt.utils.common import ( from sglang.srt.utils.common import (
direct_register_custom_op,
get_bool_env_var, get_bool_env_var,
is_cuda, is_cuda,
is_sm120_supported, is_sm120_supported,
@@ -103,6 +102,7 @@ except ImportError:
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@torch.library.custom_op("sglang::fp4_gemm", mutates_args=())
def _sglang_fp4_gemm( def _sglang_fp4_gemm(
input: torch.Tensor, input: torch.Tensor,
weight: 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) return fp4_gemm(input, weight, input_sf, weight_sf, alpha, out_dtype)
@torch.library.register_fake("sglang::fp4_gemm")
def _sglang_fp4_gemm_fake( def _sglang_fp4_gemm_fake(
input, input,
weight, weight,
@@ -135,14 +136,6 @@ def _sglang_fp4_gemm_fake(
return input.new_empty((M, N), dtype=out_dtype) 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): if is_cuda() and (not is_sm120_supported()) and (fp4_quantize is not None):
@register_fake_if_exists("sgl_kernel::scaled_fp4_quant") @register_fake_if_exists("sgl_kernel::scaled_fp4_quant")