diff --git a/python/sglang/srt/compilation/piecewise_context_manager.py b/python/sglang/srt/compilation/piecewise_context_manager.py index 55884c0d7..711b4f487 100644 --- a/python/sglang/srt/compilation/piecewise_context_manager.py +++ b/python/sglang/srt/compilation/piecewise_context_manager.py @@ -17,6 +17,9 @@ class ForwardContext: def set_attention_layers(self, layers: List[Any]): self.attention_layers = layers + def set_quant_config(self, quant_config: Any): + self.quant_config = quant_config + _forward_context: Optional[ForwardContext] = None @@ -28,11 +31,14 @@ def get_forward_context() -> Optional[ForwardContext]: @contextmanager -def set_forward_context(forward_batch: ForwardBatch, attention_layers: List[Any]): +def set_forward_context( + forward_batch: ForwardBatch, attention_layers: List[Any], quant_config: Any +): global _forward_context _forward_context = ForwardContext() _forward_context.set_forward_batch(forward_batch) _forward_context.set_attention_layers(attention_layers) + _forward_context.set_quant_config(quant_config) try: yield finally: diff --git a/python/sglang/srt/layers/quantization/awq.py b/python/sglang/srt/layers/quantization/awq.py index 437b20989..9e85a4109 100644 --- a/python/sglang/srt/layers/quantization/awq.py +++ b/python/sglang/srt/layers/quantization/awq.py @@ -954,3 +954,25 @@ class AWQMoEAscendMethod(AWQMoEMethod): use_wna16=True, ) return StandardCombineInput(hidden_states=output) + + +# Register fake implementations for torch.compile support +if _is_cuda: + + @torch.library.register_fake("sgl_kernel::awq_dequantize") + def _( + qweight, + scales, + qzeros, + ch_axis, + group_size, + num_bits, + ): + out_shape = qweight.shape[:-1] + (qweight.shape[-1] * 32 // num_bits,) + return qweight.new_empty(out_shape, dtype=scales.dtype) + + @torch.library.register_fake("sgl_kernel::awq_marlin_repack") + def _(b_q_weight, size_k, size_n, num_bits): + return b_q_weight.new_empty( + (size_k // 16, size_n * (num_bits // 2)), dtype=b_q_weight.dtype + ) diff --git a/python/sglang/srt/layers/quantization/gptq.py b/python/sglang/srt/layers/quantization/gptq.py index be28f07f8..c5f7acf25 100644 --- a/python/sglang/srt/layers/quantization/gptq.py +++ b/python/sglang/srt/layers/quantization/gptq.py @@ -1094,3 +1094,50 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase): is_k_full=self.is_k_full, ).to(orig_dtype) return StandardCombineInput(hidden_states=output) + + +# Register fake implementations for torch.compile support +if _is_cuda: + + @torch.library.register_fake("sgl_kernel::gptq_gemm") + def _(a, b_q_weight, b_gptq_qzeros, b_gptq_scales, b_g_idx, use_shuffle, bit): + return a.new_empty((a.shape[0], b_q_weight.shape[-1]), dtype=a.dtype) + + @torch.library.register_fake("sgl_kernel::gptq_marlin_repack") + def _(b_q_weight, perm, size_k, size_n, num_bits): + return b_q_weight.new_empty( + (size_k // 16, size_n * (num_bits // 2)), dtype=b_q_weight.dtype + ) + + @torch.library.register_fake("sgl_kernel::gptq_shuffle") + def _(q_weight, q_perm, bit): + return + + @torch.library.register_fake("sgl_kernel::moe_wna16_marlin_gemm") + def _( + a, + c, + b_q_weight, + b_scales, + b_zeros, + g_idx, + perm, + workspace, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + topk_weights, + moe_block_size, + top_k, + mul_topk_weights, + is_ep, + b_q_type_id, + size_m, + size_n, + size_k, + is_k_full, + use_atomic_add, + use_fp32_reduce, + is_zp_float, + ): + return c diff --git a/python/sglang/srt/layers/quantization/marlin_utils.py b/python/sglang/srt/layers/quantization/marlin_utils.py index 9a521d943..96b2efd25 100644 --- a/python/sglang/srt/layers/quantization/marlin_utils.py +++ b/python/sglang/srt/layers/quantization/marlin_utils.py @@ -31,11 +31,15 @@ if TYPE_CHECKING: from sglang.srt.layers.linear import LinearBase from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE +from sglang.srt.compilation.piecewise_context_manager import get_forward_context + try: from vllm import _custom_ops as ops except ImportError: ops = None +from sglang.srt.utils import direct_register_custom_op + _is_cuda = is_cuda() if _is_cuda: @@ -483,25 +487,44 @@ def apply_gptq_marlin_linear( dtype=input.dtype, ) - output = gptq_marlin_gemm( - reshaped_x, - None, - weight, - weight_scale, - None, - weight_zp, - g_idx, - g_idx_sort_indices, - workspace, - wtype, - size_m=reshaped_x.shape[0], - size_n=output_size_per_partition, - size_k=input_size_per_partition, - is_k_full=is_k_full, - use_atomic_add=use_atomic_add, - use_fp32_reduce=use_fp32_reduce, - is_zp_float=False, - ) + forward_context = get_forward_context() + if forward_context is None: + output = gptq_marlin_gemm( + reshaped_x, + None, + weight, + weight_scale, + None, + weight_zp, + g_idx, + g_idx_sort_indices, + workspace, + wtype, + size_m=reshaped_x.shape[0], + size_n=output_size_per_partition, + size_k=input_size_per_partition, + is_k_full=is_k_full, + use_atomic_add=use_atomic_add, + use_fp32_reduce=use_fp32_reduce, + is_zp_float=False, + ) + else: + output = torch.ops.sglang.unified_apply_gptq_marlin_gemm_with_wtype( + input=reshaped_x, + weight=weight, + weight_scale=weight_scale, + weight_zp=weight_zp, + g_idx=g_idx, + g_idx_sort_indices=g_idx_sort_indices, + workspace=workspace, + wtype_id=wtype.id, + output_size_per_partition=output_size_per_partition, + input_size_per_partition=input_size_per_partition, + is_k_full=is_k_full, + use_atomic_add=use_atomic_add, + use_fp32_reduce=use_fp32_reduce, + is_zp_float=False, + ) if bias is not None: output.add_(bias) # In-place add @@ -534,24 +557,41 @@ def apply_awq_marlin_linear( dtype=input.dtype, ) - output = gptq_marlin_gemm( - reshaped_x, - None, - weight, - weight_scale, - None, - weight_zp, - g_idx, - g_idx_sort_indices, - workspace, - quant_type, - size_m=reshaped_x.shape[0], - size_n=output_size_per_partition, - size_k=input_size_per_partition, - use_atomic_add=use_atomic_add, - use_fp32_reduce=use_fp32_reduce, - is_zp_float=False, - ) + forward_context = get_forward_context() + if forward_context is None: + output = gptq_marlin_gemm( + reshaped_x, + None, + weight, + weight_scale, + None, + weight_zp, + g_idx, + g_idx_sort_indices, + workspace, + quant_type, + size_m=reshaped_x.shape[0], + size_n=output_size_per_partition, + size_k=input_size_per_partition, + use_atomic_add=use_atomic_add, + use_fp32_reduce=use_fp32_reduce, + is_zp_float=False, + ) + else: + output = torch.ops.sglang.unified_apply_gptq_marlin_gemm( + input=reshaped_x, + weight=weight, + weight_scale=weight_scale, + weight_zp=weight_zp, + g_idx=g_idx, + g_idx_sort_indices=g_idx_sort_indices, + workspace=workspace, + output_size_per_partition=output_size_per_partition, + input_size_per_partition=input_size_per_partition, + use_atomic_add=use_atomic_add, + use_fp32_reduce=use_fp32_reduce, + is_zp_float=False, + ) if bias is not None: output.add_(bias) # In-place add @@ -818,3 +858,140 @@ class MarlinLinearMethod(LinearMethodBase): output.add_(bias) # In-place add return output + + +def unified_apply_gptq_marlin_gemm( + input: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + weight_zp: torch.Tensor, + g_idx: torch.Tensor, + g_idx_sort_indices: torch.Tensor, + workspace: torch.Tensor, + output_size_per_partition: int, + input_size_per_partition: int, + use_atomic_add: bool, + use_fp32_reduce: bool, + is_zp_float: bool, +) -> torch.Tensor: + quant_config = get_forward_context().quant_config + quant_type = quant_config.quant_type + return gptq_marlin_gemm( + input, + None, + weight, + weight_scale, + None, + weight_zp, + g_idx, + g_idx_sort_indices, + workspace, + quant_type, + size_m=input.shape[0], + size_n=output_size_per_partition, + size_k=input_size_per_partition, + use_atomic_add=use_atomic_add, + use_fp32_reduce=use_fp32_reduce, + is_zp_float=is_zp_float, + ) + + +def fake_unified_apply_gptq_marlin_gemm( + input: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + weight_zp: torch.Tensor, + g_idx: torch.Tensor, + g_idx_sort_indices: torch.Tensor, + workspace: torch.Tensor, + output_size_per_partition: int, + input_size_per_partition: int, + use_atomic_add: bool, + use_fp32_reduce: bool, + is_zp_float: bool, +) -> torch.Tensor: + return input.new_empty( + (input.shape[0], output_size_per_partition), dtype=input.dtype + ) + + +direct_register_custom_op( + op_name="unified_apply_gptq_marlin_gemm", + op_func=unified_apply_gptq_marlin_gemm, + mutates_args=[], + fake_impl=fake_unified_apply_gptq_marlin_gemm, +) + + +def unified_apply_gptq_marlin_gemm_with_wtype( + input: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + weight_zp: torch.Tensor, + g_idx: torch.Tensor, + g_idx_sort_indices: torch.Tensor, + workspace: torch.Tensor, + wtype_id: int, + output_size_per_partition: int, + input_size_per_partition: int, + is_k_full: bool, + use_atomic_add: bool, + use_fp32_reduce: bool, + is_zp_float: bool, +) -> torch.Tensor: + # Reconstruct ScalarType from id + wtype = None + for attr_name in dir(scalar_types): + if not attr_name.startswith("_"): + st = getattr(scalar_types, attr_name) + if hasattr(st, "id") and st.id == wtype_id: + wtype = st + break + return gptq_marlin_gemm( + input, + None, + weight, + weight_scale, + None, + weight_zp, + g_idx, + g_idx_sort_indices, + workspace, + wtype, + size_m=input.shape[0], + size_n=output_size_per_partition, + size_k=input_size_per_partition, + is_k_full=is_k_full, + use_atomic_add=use_atomic_add, + use_fp32_reduce=use_fp32_reduce, + is_zp_float=is_zp_float, + ) + + +def fake_unified_apply_gptq_marlin_gemm_with_wtype( + input: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + weight_zp: torch.Tensor, + g_idx: torch.Tensor, + g_idx_sort_indices: torch.Tensor, + workspace: torch.Tensor, + wtype_id: int, + output_size_per_partition: int, + input_size_per_partition: int, + is_k_full: bool, + use_atomic_add: bool, + use_fp32_reduce: bool, + is_zp_float: bool, +) -> torch.Tensor: + return input.new_empty( + (input.shape[0], output_size_per_partition), dtype=input.dtype + ) + + +direct_register_custom_op( + op_name="unified_apply_gptq_marlin_gemm_with_wtype", + op_func=unified_apply_gptq_marlin_gemm_with_wtype, + mutates_args=[], + fake_impl=fake_unified_apply_gptq_marlin_gemm_with_wtype, +) diff --git a/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py b/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py index 2a8b4ae21..ae6fbc20e 100644 --- a/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py @@ -163,6 +163,7 @@ class PiecewiseCudaGraphRunner: self.model_runner.server_args.piecewise_cuda_graph_tokens, self.model_runner.server_args.piecewise_cuda_graph_compiler, ) + self.quant_config = getattr(self.model_runner.model, "quant_config", None) # Batch sizes to capture self.capture_num_tokens = self.compile_config.get_capture_sizes() @@ -263,7 +264,9 @@ class PiecewiseCudaGraphRunner: # Attention backend self.model_runner.attn_backend.init_forward_metadata(forward_batch) - with set_forward_context(forward_batch, self.attention_layers): + with set_forward_context( + forward_batch, self.attention_layers, self.quant_config + ): _ = self.model_runner.model.forward( forward_batch.input_ids, forward_batch.positions, @@ -405,7 +408,9 @@ class PiecewiseCudaGraphRunner: set_is_extend_in_batch(False) kwargs = {} - with set_forward_context(forward_batch, self.attention_layers): + with set_forward_context( + forward_batch, self.attention_layers, self.quant_config + ): self.model_runner.model.forward( forward_batch.input_ids, forward_batch.positions, @@ -503,7 +508,9 @@ class PiecewiseCudaGraphRunner: self.model_runner.attn_backend.init_forward_metadata(forward_batch) static_forward_batch = self.replay_prepare(forward_batch, **kwargs) # Replay - with set_forward_context(static_forward_batch, self.attention_layers): + with set_forward_context( + static_forward_batch, self.attention_layers, self.quant_config + ): with set_compiled(True): output = self.model_runner.model.forward( static_forward_batch.input_ids, diff --git a/test/srt/run_suite.py b/test/srt/run_suite.py index 69698ec25..f980f8003 100644 --- a/test/srt/run_suite.py +++ b/test/srt/run_suite.py @@ -103,7 +103,7 @@ suites = { TestFile("test_original_logprobs.py", 41), TestFile("test_page_size.py", 60), TestFile("test_penalty.py", 82), - TestFile("test_piecewise_cuda_graph.py", 300), + TestFile("test_piecewise_cuda_graph.py", 450), TestFile("test_priority_scheduling.py", 130), TestFile("test_pytorch_sampling_backend.py", 66), TestFile("test_radix_attention.py", 105), diff --git a/test/srt/test_piecewise_cuda_graph.py b/test/srt/test_piecewise_cuda_graph.py index 51b92e93f..173e8b56f 100644 --- a/test/srt/test_piecewise_cuda_graph.py +++ b/test/srt/test_piecewise_cuda_graph.py @@ -132,5 +132,43 @@ class TestPiecewiseCudaGraphDeepSeek(CustomTestCase): self.assertGreater(metrics["accuracy"], 0.62) +class TestPiecewiseCudaGraphAWQ(CustomTestCase): + """Test piecewise CUDA graph with AWQ quantized model""" + + @classmethod + def setUpClass(cls): + cls.model = "Qwen/QwQ-32B-AWQ" + cls.base_url = DEFAULT_URL_FOR_TEST + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=["--enable-piecewise-cuda-graph"], + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + def test_mgsm_accuracy(self): + """Test MGSM accuracy with AWQ model""" + num_examples = 1319 + + args = SimpleNamespace( + base_url=self.base_url, + model=self.model, + eval_name="mgsm_en", + num_examples=num_examples, + num_threads=min(num_examples, 1024), + ) + + metrics = run_eval(args) + print(f"MGSM Accuracy: {metrics['score']:.3f}") + print(f"Output throughput: {metrics.get('throughput', 'N/A')} token/s") + + # Expected accuracy: 0.680, allow some variance + self.assertGreaterEqual(metrics["score"], 0.65) + + if __name__ == "__main__": unittest.main()