Quickfix for global context (#444)

Author: rebel-yskim

vllm_rbln/v1/sample/rbln_sampler.py

@@ -14,6 +14,14 @@
 # isort: off
 import torch
 import torch.nn as nn
+
+try:
+    import torch.rbln
+
+    has_torch_rbln = True
+except ImportError:
+    has_torch_rbln = False
+
 from vllm_rbln.logger import init_logger
 from vllm.v1.sample.metadata import SamplingMetadata
 from vllm.v1.sample.sampler import Sampler as VLLMSampler
@@ -151,6 +159,11 @@ def __init__(
         }
         if envs.VLLM_RBLN_COMPILE_STRICT_MODE:
             options["mode"] = "strict"
+
+        if has_torch_rbln:
+            options["use_global_ctx"] = True
+            options["global_device_id"] = 0
+
         self._compiled_rbln_topk_topp_sampler = torch.compile(
             rbln_top_k_top_p_sample,
             dynamic=False,

vllm_rbln/v1/worker/rbln_model_runner.py

@@ -1351,7 +1351,7 @@ def _compile_model(self, model):
             )
             options["cache_dir"] = os.path.join(envs.VLLM_CACHE_ROOT, "rbln")
 
-        if envs.VLLM_RBLN_AUTO_PORT and has_torch_rbln:
+        if has_torch_rbln:
             options["use_global_ctx"] = True
             # TODO(yunseong.kim): use device_id from current platform
             # when vllm-rbln supports it