Fix: Add None check in step_with_batch_queue for async scheduling

[ricky-chaoju]

Summary

Fix AttributeError: 'NoneType' object has no attribute 'sampled_token_ids' crash in V1 engine when async scheduling is enabled (default).

Root Cause

In step_with_batch_queue() method, execute_model() can return None to signal async scheduling, but there was no None check before passing model_output to scheduler.update_from_output().

The step() method already handles this correctly (line 364-366), but step_with_batch_queue() was missing the same check.

Fix

Added None check in step_with_batch_queue() to call sample_tokens() when model_output is None, mirroring the existing pattern in step().

Environment

• GPU: NVIDIA GB10 (SM 12.1 / Blackwell)
• PyTorch: 2.9.1+cu130

Relates to: #31588

[gemini-code-assist]

Code Review

This pull request correctly addresses a critical AttributeError that occurs when using asynchronous scheduling. The root cause, a missing None check for the result of execute_model in step_with_batch_queue, has been accurately identified. The fix, which involves adding a None check and calling sample_tokens to generate a valid model output, mirrors the existing logic in the step method and is a robust solution to prevent the crash. The implementation is clean, concise, and directly resolves the reported issue.

[njhill]

Thanks @rickychen-infinirc. But this should only ever happen as a secondary error after the execute_model method already raised an exception. Please share the whole log, you should see an earlier root cause error/traceback.

I agree that this is still a bug though since the secondary error is misleading. I've opened #31611 to address this.

[njhill]

Closing now that #31611 is merged.