[PieceWise CUDA Graph] Support awq/gptq model in piecewise cudagraph (#12518)

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: