fix vllm

acestep/gradio_ui.py

@@ -160,17 +160,18 @@ def create_generation_section(handler) -> dict:
         
         # Service Configuration
         with gr.Accordion("🔧 Service Configuration", open=True) as service_config_accordion:
-            with gr.Row():
-                with gr.Column(scale=2):
+            # Dropdown options section - all dropdowns grouped together
+            with gr.Row(equal_height=True):
+                with gr.Column(scale=4):
                     checkpoint_dropdown = gr.Dropdown(
                         label="Checkpoint File",
                         choices=handler.get_available_checkpoints(),
                         value=None,
                         info="Select a trained model checkpoint file (full path or filename)"
                     )
-                with gr.Column(scale=1):
+                with gr.Column(scale=1, min_width=90):
                     refresh_btn = gr.Button("🔄 Refresh", size="sm")
-                    
+            
             with gr.Row():
                 # Get available acestep-v15- model list
                 available_models = handler.get_available_acestep_v15_models()
@@ -200,13 +201,20 @@ def create_generation_section(handler) -> dict:
                     value=default_lm_model,
                     info="Select the 5Hz LM model checkpoint (auto-scanned from checkpoints)"
                 )
+                backend_dropdown = gr.Dropdown(
+                    choices=["vllm", "pt"],
+                    value="vllm",
+                    label="5Hz LM Backend",
+                    info="Select backend for 5Hz LM: vllm (faster) or pt (PyTorch, more compatible)"
+                )
+            
+            # Checkbox options section - all checkboxes grouped together
+            with gr.Row():
                 init_llm_checkbox = gr.Checkbox(
                     label="Initialize 5Hz LM",
                     value=False,
                     info="Check to initialize 5Hz LM during service initialization",
                 )
-            
-            with gr.Row():
                 # Auto-detect flash attention availability
                 flash_attn_available = handler.is_flash_attention_available()
                 use_flash_attention_checkbox = gr.Checkbox(
@@ -223,7 +231,7 @@ def create_generation_section(handler) -> dict:
                 offload_dit_to_cpu_checkbox = gr.Checkbox(
                     label="Offload DiT to CPU",
                     value=False,
-                    info="Offload DiT model to CPU when not in use (only effective if Offload to CPU is checked)"
+                    info="Offload DiT to CPU (needs Offload to CPU)"
                 )
             
             init_btn = gr.Button("Initialize Service", variant="primary", size="lg")
@@ -319,10 +327,29 @@ def create_generation_section(handler) -> dict:
                             maximum=2.0,
                             value=0.7,
                             step=0.1,
-                            scale=2,
-                            info="Temperature for 5Hz LM sampling"
+                            scale=1,
+                            info="Temperature for 5Hz LM sampling (higher = more random, lower = more deterministic)"
+                        )
+                        lm_cfg_scale = gr.Slider(
+                            label="CFG Scale",
+                            minimum=1.0,
+                            maximum=3.0,
+                            value=1.0,
+                            step=0.1,
+                            scale=1,
+                            info="Classifier-Free Guidance scale for 5Hz LM (1.0 = no CFG, higher = stronger guidance)"
                         )
                     
+                    # Negative prompt for CFG (only visible when LM initialized and cfg_scale > 1)
+                    lm_negative_prompt = gr.Textbox(
+                        label="Negative Prompt",
+                        value="NO USER INPUT",
+                        placeholder="Enter negative prompt for CFG (default: NO USER INPUT)",
+                        visible=False,
+                        info="Negative prompt used for Classifier-Free Guidance when CFG Scale > 1.0",
+                        lines=2
+                    )
+                    
                     # Repainting controls
                     with gr.Group(visible=False) as repainting_group:
                         gr.HTML("<h5>🎨 Repainting Controls (seconds) </h5>")
@@ -495,6 +522,7 @@ def create_generation_section(handler) -> dict:
         "init_status": init_status,
         "lm_model_path": lm_model_path,
         "init_llm_checkbox": init_llm_checkbox,
+        "backend_dropdown": backend_dropdown,
         "use_flash_attention_checkbox": use_flash_attention_checkbox,
         "offload_to_cpu_checkbox": offload_to_cpu_checkbox,
         "offload_dit_to_cpu_checkbox": offload_dit_to_cpu_checkbox,
@@ -510,6 +538,8 @@ def create_generation_section(handler) -> dict:
         "use_5hz_lm_row": use_5hz_lm_row,
         "use_5hz_lm_btn": use_5hz_lm_btn,
         "lm_temperature": lm_temperature,
+        "lm_cfg_scale": lm_cfg_scale,
+        "lm_negative_prompt": lm_negative_prompt,
         "repainting_group": repainting_group,
         "repainting_start": repainting_start,
         "repainting_end": repainting_end,
@@ -666,11 +696,12 @@ def setup_event_handlers(demo, handler, dataset_section, generation_section, res
     )
     
     # Service initialization
-    def init_service_wrapper(checkpoint, config_path, device, init_llm, lm_model_path, use_flash_attention, offload_to_cpu, offload_dit_to_cpu):
+    def init_service_wrapper(checkpoint, config_path, device, init_llm, lm_model_path, backend, use_flash_attention, offload_to_cpu, offload_dit_to_cpu):
         """Wrapper for service initialization, returns status and button state"""
         status, enable = handler.initialize_service(
             checkpoint, config_path, device, init_llm, lm_model_path, 
-            use_flash_attention, compile_model=False, 
+            backend=backend,
+            use_flash_attention=use_flash_attention, compile_model=False, 
             offload_to_cpu=offload_to_cpu, offload_dit_to_cpu=offload_dit_to_cpu
         )
         return status, gr.update(interactive=enable)
@@ -683,6 +714,7 @@ def setup_event_handlers(demo, handler, dataset_section, generation_section, res
             generation_section["device"],
             generation_section["init_llm_checkbox"],
             generation_section["lm_model_path"],
+            generation_section["backend_dropdown"],
             generation_section["use_flash_attention_checkbox"],
             generation_section["offload_to_cpu_checkbox"],
             generation_section["offload_dit_to_cpu_checkbox"],
@@ -690,6 +722,30 @@ def setup_event_handlers(demo, handler, dataset_section, generation_section, res
         outputs=[generation_section["init_status"], generation_section["generate_btn"]]
     )
     
+    # Update negative prompt visibility based on LM initialization and CFG scale
+    def update_negative_prompt_visibility(init_status, cfg_scale):
+        """Update negative prompt visibility: show only if LM initialized and cfg_scale > 1"""
+        # Check if LM is initialized by looking for "5Hz LM backend:" in status
+        lm_initialized = init_status is not None and "5Hz LM backend:" in str(init_status)
+        # Check if cfg_scale > 1
+        cfg_enabled = cfg_scale is not None and float(cfg_scale) > 1.0
+        # Show only if both conditions are met
+        return gr.update(visible=lm_initialized and cfg_enabled)
+    
+    # Update visibility when init_status changes
+    generation_section["init_status"].change(
+        fn=update_negative_prompt_visibility,
+        inputs=[generation_section["init_status"], generation_section["lm_cfg_scale"]],
+        outputs=[generation_section["lm_negative_prompt"]]
+    )
+    
+    # Update visibility when cfg_scale changes
+    generation_section["lm_cfg_scale"].change(
+        fn=update_negative_prompt_visibility,
+        inputs=[generation_section["init_status"], generation_section["lm_cfg_scale"]],
+        outputs=[generation_section["lm_negative_prompt"]]
+    )
+    
     # Generation with progress bar
     def generate_with_progress(
         captions, lyrics, bpm, key_scale, time_signature, vocal_language,
@@ -762,9 +818,9 @@ def setup_event_handlers(demo, handler, dataset_section, generation_section, res
     )
     
     # 5Hz LM generation (simplified version, can be extended as needed)
-    def generate_lm_hints_wrapper(caption, lyrics, temperature):
+    def generate_lm_hints_wrapper(caption, lyrics, temperature, cfg_scale, negative_prompt):
         """Wrapper for 5Hz LM generation"""
-        metadata, audio_codes, status = handler.generate_with_5hz_lm(caption, lyrics, temperature)
+        metadata, audio_codes, status = handler.generate_with_5hz_lm(caption, lyrics, temperature, cfg_scale, negative_prompt)
         
         # Extract metadata values and map to UI fields
         # Handle bpm
@@ -801,7 +857,9 @@ def setup_event_handlers(demo, handler, dataset_section, generation_section, res
         inputs=[
             generation_section["captions"],
             generation_section["lyrics"],
-            generation_section["lm_temperature"]
+            generation_section["lm_temperature"],
+            generation_section["lm_cfg_scale"],
+            generation_section["lm_negative_prompt"]
         ],
         outputs=[
             generation_section["text2music_audio_code_string"],

acestep/handler.py

@@ -151,6 +151,7 @@ class AceStepHandler:
         device: str = "auto",
         init_llm: bool = False,
         lm_model_path: str = "acestep-5Hz-lm-0.6B",
+        backend: str = "vllm",
         use_flash_attention: bool = False,
         compile_model: bool = False,
         offload_to_cpu: bool = False,
@@ -165,6 +166,7 @@ class AceStepHandler:
             device: Device type
             init_llm: Whether to initialize 5Hz LM model
             lm_model_path: 5Hz LM model path
+            backend: Backend for 5Hz LM model ("vllm" or "pt")
             use_flash_attention: Whether to use flash attention (requires flash_attn package)
             compile_model: Whether to use torch.compile to optimize the model
             offload_to_cpu: Whether to offload models to CPU when not in use
@@ -285,20 +287,20 @@ class AceStepHandler:
                 if os.path.exists(full_lm_model_path):
                     logger.info("loading 5Hz LM tokenizer...")
                     start_time = time.time()
-                    llm_tokenizer = deepcopy(self.text_tokenizer)
-                    max_audio_length = 2**16 - 1
-                    semantic_tokens = [f"<|audio_code_{i}|>" for i in range(max_audio_length)]
-                    # 217204
-                    llm_tokenizer.add_special_tokens({"additional_special_tokens": semantic_tokens})
+                    llm_tokenizer = AutoTokenizer.from_pretrained(full_lm_model_path, use_fast=True)
                     logger.info(f"5Hz LM tokenizer loaded successfully in {time.time() - start_time:.2f} seconds")
                     self.llm_tokenizer = llm_tokenizer
-                    if device == "cuda":
+                    
+                    # Initialize based on user-selected backend
+                    if backend == "vllm":
+                        # Try to initialize with vllm
                         status_msg = self._initialize_5hz_lm_vllm(full_lm_model_path)
                         logger.info(f"5Hz LM status message: {status_msg}")
                         # Check if initialization failed (status_msg starts with ❌)
                         if status_msg.startswith("❌"):
                             # vllm initialization failed, fallback to PyTorch
                             if not self.llm_initialized:
+                                logger.warning("vllm initialization failed, falling back to PyTorch backend")
                                 try:
                                     self.llm = AutoModelForCausalLM.from_pretrained(full_lm_model_path, trust_remote_code=True)
                                     if not self.offload_to_cpu:
@@ -308,15 +310,14 @@ class AceStepHandler:
                                     self.llm.eval()
                                     self.llm_backend = "pt"
                                     self.llm_initialized = True
-                                    logger.info("5Hz LM initialized successfully on CUDA device using Transformers backend")
+                                    logger.info("5Hz LM initialized successfully using PyTorch backend (fallback)")
                                 except Exception as e:
                                     return f"❌ Error initializing 5Hz LM: {str(e)}\n\nTraceback:\n{traceback.format_exc()}", False
                         # If vllm initialization succeeded, self.llm_initialized should already be True
                     else:
-                        # For CPU or other devices, use PyTorch backend
+                        # Use PyTorch backend (pt)
                         try:
                             self.llm = AutoModelForCausalLM.from_pretrained(full_lm_model_path, trust_remote_code=True)
-                            self.llm_tokenizer = AutoTokenizer.from_pretrained(full_lm_model_path, use_fast=True, trust_remote_code=True)
                             if not self.offload_to_cpu:
                                 self.llm = self.llm.to(device).to(self.dtype)
                             else:
@@ -324,7 +325,7 @@ class AceStepHandler:
                             self.llm.eval()
                             self.llm_backend = "pt"
                             self.llm_initialized = True
-                            logger.info("5Hz LM initialized successfully on non-CUDA device using Transformers backend")
+                            logger.info(f"5Hz LM initialized successfully using PyTorch backend on {device}")
                         except Exception as e:
                             return f"❌ Error initializing 5Hz LM: {str(e)}\n\nTraceback:\n{traceback.format_exc()}", False
                             
@@ -340,7 +341,9 @@ class AceStepHandler:
             status_msg += f"VAE: {vae_checkpoint_path}\n"
             status_msg += f"Text encoder: {text_encoder_path}\n"
             if init_llm and hasattr(self, 'llm') and self.llm is not None:
+                backend_info = getattr(self, 'llm_backend', 'unknown')
                 status_msg += f"5Hz LM model: {os.path.join(checkpoint_dir, lm_model_path)}\n"
+                status_msg += f"5Hz LM backend: {backend_info}\n"
             else:
                 status_msg += f"5Hz LM model: Not loaded (checkbox not selected)\n"
             status_msg += f"Dtype: {self.dtype}\n"
@@ -494,9 +497,9 @@ class AceStepHandler:
                 max_ratio=0.9
             )
             if low_gpu_memory_mode:
-                self.max_model_len = 1024
-            else:
                 self.max_model_len = 2048
+            else:
+                self.max_model_len = 4096
             
             logger.info(f"Initializing 5Hz LM with model: {model_path}, enforce_eager: False, tensor_parallel_size: 1, max_model_len: {self.max_model_len}, gpu_memory_utilization: {gpu_memory_utilization}")
             start_time = time.time()
@@ -506,9 +509,9 @@ class AceStepHandler:
                 tensor_parallel_size=1,
                 max_model_len=self.max_model_len,
                 gpu_memory_utilization=gpu_memory_utilization,
+                tokenizer=self.llm_tokenizer,
             )
             logger.info(f"5Hz LM initialized successfully in {time.time() - start_time:.2f} seconds")
-            self.llm.tokenizer = self.llm_tokenizer
             self.llm_initialized = True
             self.llm_backend = "vllm"
             return f"✅ 5Hz LM initialized successfully\nModel: {model_path}\nDevice: {device_name}\nGPU Memory Utilization: {gpu_memory_utilization:.2f}"
@@ -518,7 +521,7 @@ class AceStepHandler:
             error_msg = f"❌ Error initializing 5Hz LM: {str(e)}\n\nTraceback:\n{traceback.format_exc()}"
             return error_msg
 
-    def generate_with_5hz_lm_vllm(self, caption: str, lyrics: str, temperature: float = 0.6) -> Tuple[Dict[str, Any], str, str]:
+    def generate_with_5hz_lm_vllm(self, caption: str, lyrics: str, temperature: float = 0.6, cfg_scale: float = 1.0, negative_prompt: str = "NO USER INPUT") -> Tuple[Dict[str, Any], str, str]:
         try:
             from nanovllm import SamplingParams
             
@@ -534,35 +537,41 @@ class AceStepHandler:
             )
             logger.debug(f"[debug] formatted_prompt: {formatted_prompt}")
             
-            sampling_params = SamplingParams(max_tokens=self.max_model_len, temperature=temperature)
-            outputs = self.llm.generate([formatted_prompt], sampling_params)
+            sampling_params = SamplingParams(max_tokens=self.max_model_len-64, temperature=temperature, cfg_scale=cfg_scale)
+            # Use CFG if cfg_scale > 1.0
+            if cfg_scale > 1.0:
+                # Build unconditional prompt (user input replaced with "NO USER INPUT")
+                formatted_unconditional_prompt = self.lm_tokenizer.apply_chat_template(
+                    [
+                        {"role": "system", "content": "# Instruction\nGenerate audio semantic tokens based on the given conditions:\n\n"},
+                        {"role": "user", "content": negative_prompt}
+                    ],
+                    tokenize=False,
+                    add_generation_prompt=True,
+                )
+                outputs = self.llm.generate(
+                    [formatted_prompt], 
+                    sampling_params,
+                    unconditional_prompts=[formatted_unconditional_prompt]
+                )
+            else:
+                outputs = self.lm_model.generate([formatted_prompt], sampling_params)
+            # Extract text from output - handle different output formats
             if isinstance(outputs, list) and len(outputs) > 0:
                 if hasattr(outputs[0], 'outputs') and len(outputs[0].outputs) > 0:
                     output_text = outputs[0].outputs[0].text
                 elif hasattr(outputs[0], 'text'):
                     output_text = outputs[0].text
+                elif isinstance(outputs[0], dict) and 'text' in outputs[0]:
+                    output_text = outputs[0]['text']
                 else:
-                    # Transformers generation
-                    inputs = self.llm_tokenizer(formatted_prompt, return_tensors="pt").to(self.llm.device)
-                    
-                    # Generate
-                    with torch.no_grad():
-                        outputs = self.llm.generate(
-                            **inputs,
-                            max_new_tokens=3072,
-                            temperature=temperature,
-                            do_sample=True,
-                            pad_token_id=self.llm_tokenizer.pad_token_id,
-                            eos_token_id=self.llm_tokenizer.eos_token_id
-                        )
-                    
-                    # Decode
-                    generated_ids = outputs[0][inputs.input_ids.shape[1]:]
-                    output_text = self.llm_tokenizer.decode(generated_ids, skip_special_tokens=False)
-                
-                metadata, audio_codes = self.parse_lm_output(output_text)
-                codes_count = len(audio_codes.split('<|audio_code_')) - 1 if audio_codes else 0
-                return metadata, audio_codes, f"✅ Generated successfully\nOutput length: {len(output_text)} chars\nCodes count: {codes_count}"
+                    output_text = str(outputs[0])
+            else:
+                output_text = str(outputs)
+            metadata, audio_codes = self.parse_lm_output(output_text)
+            print(f"[debug]output_text: {output_text}")
+            codes_count = len(audio_codes.split('<|audio_code_')) - 1 if audio_codes else 0
+            return metadata, audio_codes, f"✅ Generated successfully\nOutput length: {len(output_text)} chars\nCodes count: {codes_count}"
             
         except Exception as e:
             error_msg = f"❌ Error generating with 5Hz LM: {str(e)}\n\nTraceback:\n{traceback.format_exc()}"
@@ -639,7 +648,7 @@ class AceStepHandler:
             error_msg = f"❌ Error generating with 5Hz LM: {str(e)}\n\nTraceback:\n{traceback.format_exc()}"
             return {}, "", error_msg
         
-    def generate_with_5hz_lm(self, caption: str, lyrics: str, temperature: float = 0.6) -> Tuple[Dict[str, Any], str, str]:
+    def generate_with_5hz_lm(self, caption: str, lyrics: str, temperature: float = 0.6, cfg_scale: float = 1.0, negative_prompt: str = "NO USER INPUT") -> Tuple[Dict[str, Any], str, str]:
         """Generate metadata and audio codes using 5Hz LM"""
         # Check if 5Hz LM is initialized
         if not hasattr(self, 'llm_initialized') or not self.llm_initialized:
@@ -656,7 +665,7 @@ class AceStepHandler:
             return {}, "", "❌ 5Hz LM backend not set. Please initialize it first."
         
         if self.llm_backend == "vllm":
-            return self.generate_with_5hz_lm_vllm(caption, lyrics, temperature)
+            return self.generate_with_5hz_lm_vllm(caption, lyrics, temperature, cfg_scale, negative_prompt)
         else:
             return self.generate_with_5hz_lm_pt(caption, lyrics, temperature)
     

acestep/third_parts/nano-vllm/nanovllm/config.py

@@ -1,35 +1,8 @@
 import os
-import socket
 from dataclasses import dataclass
 from transformers import AutoConfig
 
 
-def find_available_port(start_port: int = 2333, max_attempts: int = 100) -> int:
-    """Find an available port starting from start_port.
-    
-    Args:
-        start_port: The starting port number to check
-        max_attempts: Maximum number of ports to try
-        
-    Returns:
-        An available port number
-        
-    Raises:
-        RuntimeError: If no available port is found within max_attempts
-    """
-    for i in range(max_attempts):
-        port = start_port + i
-        try:
-            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
-                s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
-                s.bind(('localhost', port))
-                return port
-        except OSError:
-            # Port is in use, try next one
-            continue
-    raise RuntimeError(f"Could not find an available port starting from {start_port} after {max_attempts} attempts")
-
-
 @dataclass
 class Config:
     model: str
@@ -40,10 +13,9 @@ class Config:
     tensor_parallel_size: int = 1
     enforce_eager: bool = False
     hf_config: AutoConfig | None = None
-    eos: int = 151643
+    eos: int = -1
     kvcache_block_size: int = 256
     num_kvcache_blocks: int = -1
-    dist_port: int | None = None
 
     def __post_init__(self):
         assert os.path.isdir(self.model)
@@ -52,6 +24,3 @@ class Config:
         self.hf_config = AutoConfig.from_pretrained(self.model)
         self.max_model_len = min(self.max_model_len, self.hf_config.max_position_embeddings)
         assert self.max_num_batched_tokens >= self.max_model_len
-        # Auto-find available port if not specified
-        if self.dist_port is None:
-            self.dist_port = find_available_port()

acestep/third_parts/nano-vllm/nanovllm/engine/llm_engine.py

@@ -21,28 +21,6 @@ class LLMEngine:
         self.ps = []
         self.events = []
         ctx = mp.get_context("spawn")
-        
-        # Pre-validate port availability by attempting to bind to it
-        # This helps avoid race conditions when multiple LLMEngine instances start simultaneously
-        import socket
-        from nanovllm.config import find_available_port
-        max_port_retries = 10
-        for port_attempt in range(max_port_retries):
-            try:
-                # Test if port is actually available by binding to it
-                test_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-                test_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
-                test_socket.bind(('localhost', config.dist_port))
-                test_socket.close()
-                # Port is available, break
-                break
-            except OSError:
-                # Port is in use, find next available
-                if port_attempt < max_port_retries - 1:
-                    config.dist_port = find_available_port(start_port=config.dist_port + 1, max_attempts=10)
-                else:
-                    raise RuntimeError(f"Failed to find available port after {max_port_retries} attempts")
-        
         for i in range(1, config.tensor_parallel_size):
             event = ctx.Event()
             process = ctx.Process(target=ModelRunner, args=(config, i, event))
@@ -50,7 +28,12 @@ class LLMEngine:
             self.ps.append(process)
             self.events.append(event)
         self.model_runner = ModelRunner(config, 0, self.events)
-        self.tokenizer = None
+        tokenizer = kwargs.get("tokenizer", None)
+        if tokenizer is not None:
+            self.tokenizer = tokenizer
+        else:
+            self.tokenizer = AutoTokenizer.from_pretrained(config.model, use_fast=True)
+        config.eos = self.tokenizer.eos_token_id
         self.scheduler = Scheduler(config)
         atexit.register(self.exit)
 

acestep/third_parts/nano-vllm/nanovllm/engine/model_runner.py

@@ -1,17 +1,44 @@
 import pickle
-import socket
 import torch
 import torch.distributed as dist
 from multiprocessing.synchronize import Event
 from multiprocessing.shared_memory import SharedMemory
 
-from nanovllm.config import Config, find_available_port
+from nanovllm.config import Config
 from nanovllm.engine.sequence import Sequence
 from nanovllm.models.qwen3 import Qwen3ForCausalLM
 from nanovllm.layers.sampler import Sampler
 from nanovllm.utils.context import set_context, get_context, reset_context
 from nanovllm.utils.loader import load_model
 
+import socket
+
+
+def find_available_port(start_port: int = 2333, max_attempts: int = 100) -> int:
+    """Find an available port starting from start_port.
+    
+    Args:
+        start_port: The starting port number to check
+        max_attempts: Maximum number of ports to try
+        
+    Returns:
+        An available port number
+        
+    Raises:
+        RuntimeError: If no available port is found within max_attempts
+    """
+    for i in range(max_attempts):
+        port = start_port + i
+        try:
+            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+                s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+                s.bind(('localhost', port))
+                return port
+        except OSError:
+            # Port is in use, try next one
+            continue
+    raise RuntimeError(f"Could not find an available port starting from {start_port} after {max_attempts} attempts")
+
 
 class ModelRunner:
 
@@ -23,33 +50,9 @@ class ModelRunner:
         self.world_size = config.tensor_parallel_size
         self.rank = rank
         self.event = event
-
-        # Try to initialize process group with retry logic for port conflicts
-        # Only rank 0 binds to the port, so only rank 0 needs retry logic
-        dist_port = self.config.dist_port
-        max_retries = 10
-        for attempt in range(max_retries):
-            try:
-                dist.init_process_group("nccl", f"tcp://localhost:{dist_port}", world_size=self.world_size, rank=rank)
-                break
-            except RuntimeError as e:
-                if ("EADDRINUSE" in str(e) or "address already in use" in str(e).lower()) and rank == 0:
-                    # Port is in use, try next port (only for rank 0)
-                    if attempt < max_retries - 1:
-                        # Find next available port
-                        dist_port = find_available_port(start_port=dist_port + 1, max_attempts=10)
-                        self.config.dist_port = dist_port
-                        # If we had a previous failed attempt, destroy any partial process group
-                        if dist.is_initialized():
-                            try:
-                                dist.destroy_process_group()
-                            except:
-                                pass
-                    else:
-                        raise RuntimeError(f"Failed to find available port after {max_retries} attempts. Last error: {e}")
-                else:
-                    # Other error or non-rank-0 process, re-raise
-                    raise
+        dist_port = find_available_port()
+        print(f"[debug]dist_port: {dist_port}")
+        dist.init_process_group("nccl", f"tcp://localhost:{dist_port}", world_size=self.world_size, rank=rank)
         torch.cuda.set_device(rank)
         default_dtype = torch.get_default_dtype()
         torch.set_default_dtype(hf_config.torch_dtype)
@@ -144,15 +147,9 @@ class ModelRunner:
                 layer_id += 1
 
     def prepare_block_tables(self, seqs: list[Sequence]):
-        max_len = max(len(seq.block_table) for seq in seqs) if seqs else 0
-        if max_len == 0:
-            # Return empty 2D tensor with correct shape
-            return torch.zeros((len(seqs), 0), dtype=torch.int32, pin_memory=True).cuda(non_blocking=True)
+        max_len = max(len(seq.block_table) for seq in seqs)
         block_tables = [seq.block_table + [-1] * (max_len - len(seq.block_table)) for seq in seqs]
         block_tables = torch.tensor(block_tables, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True)
-        # Ensure it's 2D: if only one sequence, shape should be [1, max_len]
-        if block_tables.dim() == 1:
-            block_tables = block_tables.unsqueeze(0)
         return block_tables
 
     def prepare_prefill(self, seqs: list[Sequence]):
@@ -247,29 +244,7 @@ class ModelRunner:
             graph_vars["slot_mapping"][:bs] = context.slot_mapping
             graph_vars["context_lens"].zero_()
             graph_vars["context_lens"][:bs] = context.context_lens
-            # Handle block_tables: ensure it's 2D and size matches
-            if context.block_tables is not None and context.block_tables.numel() > 0:
-                # Ensure block_tables is 2D
-                if context.block_tables.dim() == 1:
-                    # Reshape 1D to 2D: [num_blocks] -> [1, num_blocks]
-                    block_tables_2d = context.block_tables.unsqueeze(0)
-                else:
-                    block_tables_2d = context.block_tables
-                
-                # Get dimensions
-                context_bs = block_tables_2d.size(0)
-                context_num_blocks = block_tables_2d.size(1)
-                graph_num_blocks = graph_vars["block_tables"].size(1)
-                
-                # Use minimum to avoid size mismatch
-                num_blocks_to_copy = min(context_num_blocks, graph_num_blocks)
-                actual_bs = min(bs, context_bs)
-                
-                # Copy block_tables with size matching
-                graph_vars["block_tables"][:actual_bs, :num_blocks_to_copy] = block_tables_2d[:actual_bs, :num_blocks_to_copy]
-                # Fill remaining with -1 if needed
-                if num_blocks_to_copy < graph_num_blocks:
-                    graph_vars["block_tables"][:actual_bs, num_blocks_to_copy:] = -1
+            graph_vars["block_tables"][:bs, :context.block_tables.size(1)] = context.block_tables
             graph.replay()
             return self.model.compute_logits(graph_vars["outputs"][:bs])