use logged_rbln_backend in more places

Author: rebel-jaehwang

vllm_rbln/v1/sample/rbln_sampler.py

@@ -24,6 +24,7 @@
     has_torch_rbln = False
 
 from vllm_rbln.logger import init_logger
+from vllm_rbln.torch_compile_backend import logged_rbln_backend
 from vllm.v1.sample.metadata import SamplingMetadata
 from vllm.v1.sample.sampler import Sampler as VLLMSampler
 import rebel
@@ -175,7 +176,7 @@ def __init__(
             rbln_top_k_top_p_sample,
             dynamic=False,
             fullgraph=True,
-            backend="rbln",
+            backend=logged_rbln_backend,
             options=options,
         )
         self.forward = self.forward_rbln

vllm_rbln/v1/spec_decode/eagle.py

@@ -30,6 +30,7 @@
 import vllm_rbln.rbln_envs as envs
 import vllm_rbln.utils as rbln_utils
 from vllm_rbln.logger import init_logger
+from vllm_rbln.torch_compile_backend import logged_rbln_backend
 from vllm_rbln.v1.attention.kv_cache_bindings import (
     attach_kv_cache_bindings,
     build_kv_cache_forward_context_kwargs,
@@ -515,7 +516,7 @@ def _compile_model(self, model):
 
         return torch.compile(
             model,
-            backend="rbln",
+            backend=logged_rbln_backend,
             options=copy(options),
             dynamic=False,
         )

vllm_rbln/v1/spec_decode/medusa.py

@@ -23,6 +23,7 @@
 from vllm.v1.spec_decode.medusa import MedusaProposer
 
 import vllm_rbln.rbln_envs as envs
+from vllm_rbln.torch_compile_backend import logged_rbln_backend
 
 
 class RBLNMedusaProposer(MedusaProposer):
@@ -74,7 +75,7 @@ def _compile_model(self, model: nn.Module):
 
         return torch.compile(
             model,
-            backend="rbln",
+            backend=logged_rbln_backend,
             options=copy(options),
             dynamic=False,
         )