fix: Re-enable standalone ctx (#407)

* Revert "Revert "fix(core): Sampler with `RBLN_CTX_STANDALONE` (#401)" (#405)"

This reverts commit 81f7d47.

* Revert "fix RBLN_CTX_STANDALONE=0"

This reverts commit defd60e.

---------

Co-authored-by: rebel-jonghewk <jonghewk@rebellions.in>

Author: rebel-jonghewk

vllm_rbln/platform.py

@@ -135,7 +135,7 @@ def validate_and_setup_prerequisite(cls, vllm_config: VllmConfig) -> None:
                     "RBLN_PROFILER is not supported when using vLLM model parallel "
                     "(TP, DP, EP, or PP)."
                 )
-            os.environ["RBLN_CTX_STANDALONE"] = "0"
+            os.environ["RBLN_CTX_STANDALONE"] = "1"
             os.environ["RBLN_FORCE_CCL_ASYNC"] = "1"
 
     @classmethod

vllm_rbln/v1/sample/rbln_sampler.py

@@ -128,7 +128,12 @@ def rbln_top_k_top_p_sample(
 
 
 class RBLNTopKTopPSampler(nn.Module):
-    def __init__(self, logprobs_mode: LogprobsMode = "raw_logprobs", seed: int = 42):
+    def __init__(
+        self,
+        logprobs_mode: LogprobsMode = "raw_logprobs",
+        seed: int = 42,
+        compile_context: rebel.CompileContext = None,
+    ):
         # TODO(rbln): Merge more ops to rbln context.
         #       Currently, we only have softmax in rbln context.
         super().__init__()
@@ -139,7 +144,11 @@ def __init__(self, logprobs_mode: LogprobsMode = "raw_logprobs", seed: int = 42)
         )
 
         rebel.manual_seed(seed)
-        options = {"compile_context": rebel.CompileContext()}
+        options = {
+            "compile_context": compile_context
+            if compile_context
+            else rebel.CompileContext()
+        }
         if envs.VLLM_RBLN_COMPILE_STRICT_MODE:
             options["mode"] = "strict"
         self._compiled_rbln_topk_topp_sampler = torch.compile(
@@ -175,11 +184,16 @@ def forward_rbln(
 
 
 class RBLNSampler(VLLMSampler):
-    def __init__(self, logprobs_mode: LogprobsMode = "raw_logprobs", seed: int = 42):
+    def __init__(
+        self,
+        logprobs_mode: LogprobsMode = "raw_logprobs",
+        seed: int = 42,
+        compile_context: rebel.CompileContext = None,
+    ):
         super().__init__()
         if logprobs_mode in ("raw_logprobs", "raw_logits"):
             self.topk_topp_sampler = RBLNTopKTopPSampler(
-                logprobs_mode=logprobs_mode, seed=seed
+                logprobs_mode=logprobs_mode, seed=seed, compile_context=compile_context
             )
         else:
             logger.warning_once(

vllm_rbln/v1/worker/rbln_model_runner.py

@@ -272,13 +272,18 @@ def __init__(
         else:
             self.max_encoder_len = 0
 
+        from rebel.compile_context import CompileContext
+
+        self.compile_context = CompileContext(use_weight_sharing=True)
+
         # Sampler
         self.use_rbln_sampler = envs.VLLM_RBLN_SAMPLER
         if self.use_rbln_sampler:
             logger.info("Using RBLN sampler: %s", self.use_rbln_sampler)
             sampler = RBLNSampler(
                 logprobs_mode=self.model_config.logprobs_mode,
                 seed=self.vllm_config.model_config.seed,
+                compile_context=self.compile_context,
             )
         else:
             logger.info("Using default vLLM sampler.")
@@ -3211,9 +3216,7 @@ def model_wrapper(
             # RBLN compile context to mark static address for kv cache tensor
             # if tensor is set to have static address,
             # similar to RBLN kv cache binding
-            from rebel.compile_context import CompileContext
 
-            self.compile_context = CompileContext(use_weight_sharing=True)
             compiled_graph = self._compile_model(model_wrapper)
             self.model_executable = compiled_graph