diff --git a/.github/workflows/gpu-integration.yml b/.github/workflows/gpu-integration.yml
index d7a0c6b2..21b9b53e 100644
--- a/.github/workflows/gpu-integration.yml
+++ b/.github/workflows/gpu-integration.yml
@@ -103,7 +103,7 @@ jobs:
             });
 
   integration:
-    name: Qwen3.8 Eagle3 + DSpark (GB200, 2 GPUs)
+    name: Qwen3.8 Eagle3 + DSpark + DFlash2 (GB200, 2 GPUs)
     needs: authorize
     runs-on: ${{ vars.TORCHSPEC_GPU_DISPATCHER_LABEL }}
     timeout-minutes: 720
diff --git a/.github/workflows/gpu-nightly.yml b/.github/workflows/gpu-nightly.yml
index 45acaf9f..00a97bbb 100644
--- a/.github/workflows/gpu-nightly.yml
+++ b/.github/workflows/gpu-nightly.yml
@@ -58,7 +58,7 @@ jobs:
           } >> "$GITHUB_STEP_SUMMARY"
 
   integration:
-    name: Qwen3.8 full-epoch GPU training
+    name: Qwen3.8 Eagle3 + DSpark + DFlash2 GPU training
     needs: changes
     if: needs.changes.outputs.should_run == 'true'
     runs-on: ${{ vars.TORCHSPEC_GPU_DISPATCHER_LABEL }}
diff --git a/configs/ci/qwen3_8_27b_dflash2_draft.json b/configs/ci/qwen3_8_27b_dflash2_draft.json
new file mode 100644
index 00000000..8917abe7
--- /dev/null
+++ b/configs/ci/qwen3_8_27b_dflash2_draft.json
@@ -0,0 +1,37 @@
+{
+  "architectures": ["DFlash2DraftModel"],
+  "model_type": "qwen3_dflash2",
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "hidden_size": 5120,
+  "intermediate_size": 8192,
+  "num_hidden_layers": 5,
+  "num_attention_heads": 40,
+  "num_key_value_heads": 8,
+  "head_dim": 128,
+  "vocab_size": 248320,
+  "rms_norm_eps": 1e-6,
+  "max_position_embeddings": 262144,
+  "rope_theta": 10000000.0,
+  "num_target_layers": 5,
+  "target_hidden_size": 5120,
+  "target_num_hidden_layers": 64,
+  "layer_types": [
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention"
+  ],
+  "is_causal": false,
+  "dflash_config": {
+    "block_size": 8,
+    "conv_kernel_size": 2,
+    "conv_group_size": 16,
+    "selector_rank": 256,
+    "selector_top_k": 16,
+    "target_layer_ids": [3, 15, 31, 47, 59],
+    "mask_token_id": 248044
+  },
+  "tie_word_embeddings": false
+}
diff --git a/configs/ci/qwen3_8_27b_dspark_draft.json b/configs/ci/qwen3_8_27b_dspark_draft.json
index 1148b450..1372830b 100644
--- a/configs/ci/qwen3_8_27b_dspark_draft.json
+++ b/configs/ci/qwen3_8_27b_dspark_draft.json
@@ -3,7 +3,7 @@
   "model_type": "qwen3_dspark",
   "hidden_size": 5120,
   "intermediate_size": 8192,
-  "num_hidden_layers": 1,
+  "num_hidden_layers": 5,
   "num_attention_heads": 40,
   "num_key_value_heads": 8,
   "head_dim": 128,
@@ -16,7 +16,7 @@
   "target_num_hidden_layers": 64,
   "target_layer_ids": [3, 15, 31, 47, 59],
   "mask_token_id": 248044,
-  "markov_rank": 64,
+  "markov_rank": 256,
   "markov_head_type": "vanilla",
   "enable_confidence_head": true,
   "confidence_head_with_markov": true,
diff --git a/configs/ci/vllm_qwen3_8_27b_dflash2_2gpu_smoke.yaml b/configs/ci/vllm_qwen3_8_27b_dflash2_2gpu_smoke.yaml
new file mode 100644
index 00000000..aef3d905
--- /dev/null
+++ b/configs/ci/vllm_qwen3_8_27b_dflash2_2gpu_smoke.yaml
@@ -0,0 +1,81 @@
+# Two-GPU CI lane: the same Qwen3.8-27B target on one GB200 and a
+# production-depth DFlash2 draft on the other. This exercises block diffusion,
+# grouped convolutions, candidate selection, and selector-loss training.
+model:
+  target_model_path: Qwen/Qwen3.8-27B
+  target_model_backend: vllm
+  draft_model_config: configs/ci/qwen3_8_27b_dflash2_draft.json
+  embedding_key: model.language_model.embed_tokens.weight
+  lm_head_key: lm_head.weight
+  norm_key: model.language_model.norm.weight
+  trust_remote_code: true
+
+dataset:
+  train_data_path: ../../examples/data/sample_conversations.jsonl
+  chat_template: qwen
+  prompt_key: conversations
+  shuffle_dataset: false
+  length_group_size: 1
+  min_loss_tokens: 32
+
+training:
+  attention_backend: flex_attention
+  micro_batch_size: 1
+  draft_accumulation_steps: 1
+  learning_rate: 1e-4
+  max_concurrent_batches: 1
+  max_grad_norm: 0.5
+  # Keep the complete 12,913-token record so DFlash2 is exercised after
+  # multiple vLLM chunked-prefill iterations.
+  max_seq_length: 16384
+  num_epochs: 1
+  prefetch_depth: 1
+  save_interval: 0
+  save_per_epoch: false
+  seed: 42
+  training_num_gpus_per_node: 1
+  training_num_nodes: 1
+  warmup_ratio: 0.0
+  dflash_block_size: 8
+  dflash_num_anchors: 512
+  dflash_num_target_layers: 5
+  dflash_loss_decay_gamma: 7.0
+  dflash2_selector_loss_alpha: 1.0
+
+inference:
+  inference_engine_type: vllm
+  store_last_hidden_states: false
+  inference_num_gpus: 1
+  inference_num_gpus_per_engine: 1
+  inference_num_gpus_per_node: 2
+  inference_batch_size: 1
+  inference_buffer_threshold: 2
+  max_sample_pool_size: 4
+  # These must match dflash_config.target_layer_ids in the draft config.
+  aux_hidden_states_layers: [3, 15, 31, 47, 59]
+  vllm:
+    tp_size: 1
+    mem_fraction_static: 0.55
+    extra_args:
+      enforce_eager: true
+      language_model_only: true
+      max_model_len: 16384
+      max_num_batched_tokens: 4096
+
+mooncake:
+  master_server_address: null
+  metadata_server: null
+  protocol: tcp
+  global_segment_size: 4GB
+  local_buffer_size: 1GB
+
+output_dir: ""
+cache_dir: ../../.ci-cache/qwen3-8-27b-dflash2-smoke
+model_download_dir: null
+
+logging:
+  report_to: none
+
+debug:
+  enable_perf_metrics: true
+  save_debug_train_data: null
diff --git a/configs/ci/vllm_qwen3_8_27b_dspark_2gpu_smoke.yaml b/configs/ci/vllm_qwen3_8_27b_dspark_2gpu_smoke.yaml
index 326c42bc..2cd59283 100644
--- a/configs/ci/vllm_qwen3_8_27b_dspark_2gpu_smoke.yaml
+++ b/configs/ci/vllm_qwen3_8_27b_dspark_2gpu_smoke.yaml
@@ -1,6 +1,6 @@
 # Two-GPU CI lane: the same Qwen3.8-27B target on one GB200 and a
-# one-layer DSpark draft on the other. This exercises block diffusion,
-# Markov, confidence, CE, and L1 training paths without a production-sized draft.
+# production-depth DSpark draft on the other. This exercises block diffusion,
+# Markov, confidence, CE, and L1 training paths with the standard DSpark shape.
 model:
   target_model_path: Qwen/Qwen3.8-27B
   target_model_backend: vllm
@@ -16,7 +16,7 @@ dataset:
   prompt_key: conversations
   shuffle_dataset: false
   length_group_size: 1
-  min_loss_tokens: 4
+  min_loss_tokens: 32
 
 training:
   attention_backend: flex_attention
@@ -36,8 +36,8 @@ training:
   training_num_gpus_per_node: 1
   training_num_nodes: 1
   warmup_ratio: 0.0
-  dflash_block_size: 2
-  dspark_num_anchors: 4
+  dflash_block_size: 7
+  dspark_num_anchors: 512
   dspark_num_target_layers: 5
   dspark_loss_decay_gamma: 4.0
   dspark_ce_loss_alpha: 0.1
@@ -80,5 +80,5 @@ logging:
   report_to: none
 
 debug:
-  enable_perf_metrics: false
+  enable_perf_metrics: true
   save_debug_train_data: null
diff --git a/configs/ci/vllm_qwen3_8_27b_eagle3_2gpu_smoke.yaml b/configs/ci/vllm_qwen3_8_27b_eagle3_2gpu_smoke.yaml
index 6b2d25f2..2a86eb1d 100644
--- a/configs/ci/vllm_qwen3_8_27b_eagle3_2gpu_smoke.yaml
+++ b/configs/ci/vllm_qwen3_8_27b_eagle3_2gpu_smoke.yaml
@@ -74,5 +74,5 @@ logging:
   report_to: none
 
 debug:
-  enable_perf_metrics: false
+  enable_perf_metrics: true
   save_debug_train_data: null
diff --git a/tests/test_dflash2.py b/tests/test_dflash2.py
index bf834d97..efc9c8a0 100644
--- a/tests/test_dflash2.py
+++ b/tests/test_dflash2.py
@@ -35,6 +35,7 @@
     _convert_fsdp_to_hf,
     _fixup_export_config,
 )
+from torchspec.config import load_config
 from torchspec.models.dflash import _create_dflash_mask_mod
 from torchspec.models.dflash2 import DFlash2Model
 from torchspec.models.draft.auto import AutoDraftModelConfig
@@ -171,6 +172,29 @@ def test_repository_config_dispatches_to_dflash2(self):
         with mock.patch.dict("sys.modules", {"torchspec.training.dflash2_trainer": trainer_module}):
             self.assertIs(_trainer_class_for_config(config), trainer_class)
 
+    def test_ci_config_dispatches_to_dflash2_with_matching_capture_layers(self):
+        config_path = ROOT / "configs" / "ci" / "vllm_qwen3_8_27b_dflash2_2gpu_smoke.yaml"
+        config = load_config(str(config_path))
+        draft_config = AutoDraftModelConfig.from_file(config.model.draft_model_config)
+
+        self.assertIsInstance(draft_config, DFlash2Config)
+        self.assertEqual(draft_config.architectures, ["DFlash2DraftModel"])
+        self.assertEqual(
+            list(config.inference.aux_hidden_states_layers),
+            draft_config.target_layer_ids,
+        )
+        self.assertEqual(config.training.dflash_block_size, draft_config.block_size)
+        self.assertEqual(config.training.dflash_num_anchors, 512)
+        self.assertEqual(config.dataset.min_loss_tokens, 32)
+        self.assertEqual(
+            config.training.dflash_num_target_layers,
+            draft_config.num_target_layers,
+        )
+        self.assertEqual(draft_config.num_hidden_layers, 5)
+        self.assertEqual(draft_config.target_num_hidden_layers, 64)
+        self.assertEqual(draft_config.selector_rank, 256)
+        self.assertEqual(draft_config.selector_top_k, 16)
+
     def test_legacy_dflash_dispatch_is_unchanged(self):
         config = AutoDraftModelConfig.from_dict(
             {
diff --git a/tests/test_dspark.py b/tests/test_dspark.py
index 59d4c0c1..90b2aff4 100644
--- a/tests/test_dspark.py
+++ b/tests/test_dspark.py
@@ -37,9 +37,11 @@
 
 import math
 import unittest
+from pathlib import Path
 
 import torch
 
+from torchspec.config import load_config
 from torchspec.models.draft.auto import AutoDraftModelConfig, AutoEagle3DraftModel
 from torchspec.models.draft.dflash import DFlashConfig
 from torchspec.models.draft.dspark import (
@@ -54,6 +56,7 @@
 from torchspec.models.dspark import DSparkModel
 
 CE_A, L1_A, CF_A = 0.1, 0.9, 1.0
+ROOT = Path(__file__).resolve().parents[1]
 
 
 def _make_dspark_config(
@@ -135,6 +138,23 @@ def test_subclasses_dflash_and_attrs(self):
         self.assertTrue(cfg.enable_confidence_head)
         self.assertFalse(cfg.fc_norm)
 
+    def test_ci_config_uses_standard_dspark_shape(self):
+        config_path = ROOT / "configs" / "ci" / "vllm_qwen3_8_27b_dspark_2gpu_smoke.yaml"
+        config = load_config(str(config_path))
+        draft_config = AutoDraftModelConfig.from_file(config.model.draft_model_config)
+
+        self.assertIsInstance(draft_config, DSparkConfig)
+        self.assertEqual(config.training.dflash_block_size, 7)
+        self.assertEqual(config.training.dspark_num_anchors, 512)
+        self.assertEqual(config.dataset.min_loss_tokens, 32)
+        self.assertEqual(draft_config.num_hidden_layers, 5)
+        self.assertEqual(draft_config.markov_rank, 256)
+        self.assertEqual(draft_config.target_num_hidden_layers, 64)
+        self.assertEqual(
+            list(config.inference.aux_hidden_states_layers),
+            draft_config.target_layer_ids,
+        )
+
     def test_optional_fc_norm_normalizes_each_target_layer_before_projection(self):
         cfg = _make_dspark_config(H=16, num_target_layers=3, fc_norm=True)
         model = DSparkDraftModel(cfg).to(dtype=torch.float32)
diff --git a/tests/test_env.py b/tests/test_env.py
new file mode 100644
index 00000000..54dbc214
--- /dev/null
+++ b/tests/test_env.py
@@ -0,0 +1,19 @@
+"""Tests for the default environment forwarded to TorchSpec Ray actors."""
+
+from torchspec.utils.env import get_torchspec_env_vars
+
+
+def test_vllm_v2_model_runner_is_enabled_by_default(monkeypatch):
+    monkeypatch.delenv("VLLM_USE_V2_MODEL_RUNNER", raising=False)
+
+    env = get_torchspec_env_vars()
+
+    assert env["VLLM_USE_V2_MODEL_RUNNER"] == "1"
+
+
+def test_explicit_vllm_model_runner_value_is_forwarded(monkeypatch):
+    monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0")
+
+    env = get_torchspec_env_vars()
+
+    assert env["VLLM_USE_V2_MODEL_RUNNER"] == "0"
diff --git a/tools/ci/gpu_2gpu.sbatch b/tools/ci/gpu_2gpu.sbatch
index 6c956cb8..d0f1e635 100755
--- a/tools/ci/gpu_2gpu.sbatch
+++ b/tools/ci/gpu_2gpu.sbatch
@@ -24,6 +24,7 @@ srun --gpus=2 \
     export RUNNER_TEMP=/ci-tmp
     export HF_HOME=/ci-tmp/huggingface
     export PYTHONPATH="/workspace${PYTHONPATH:+:${PYTHONPATH}}"
+    export VLLM_USE_V2_MODEL_RUNNER=1
     export TORCHSPEC_CI_ARTIFACT_DIR=/artifacts
     export TORCHSPEC_CI_MODEL_PATH="/model-cache/snapshots/${TORCHSPEC_CI_MODEL_REVISION}"
     export TORCHSPEC_CI_MODEL_CACHE=/ci-tmp/huggingface
diff --git a/tools/ci/run_2gpu_training_smoke.sh b/tools/ci/run_2gpu_training_smoke.sh
index 83053724..2daa1c99 100755
--- a/tools/ci/run_2gpu_training_smoke.sh
+++ b/tools/ci/run_2gpu_training_smoke.sh
@@ -4,6 +4,7 @@ set -euo pipefail
 repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
 eagle3_config="${TORCHSPEC_CI_EAGLE3_CONFIG:-${repo_root}/configs/ci/vllm_qwen3_8_27b_eagle3_2gpu_smoke.yaml}"
 dspark_config="${TORCHSPEC_CI_DSPARK_CONFIG:-${repo_root}/configs/ci/vllm_qwen3_8_27b_dspark_2gpu_smoke.yaml}"
+dflash2_config="${TORCHSPEC_CI_DFLASH2_CONFIG:-${repo_root}/configs/ci/vllm_qwen3_8_27b_dflash2_2gpu_smoke.yaml}"
 fixture="${TORCHSPEC_CI_FIXTURE:-${repo_root}/examples/data/sample_conversations.jsonl}"
 artifact_dir="${TORCHSPEC_CI_ARTIFACT_DIR:-${RUNNER_TEMP:-/tmp}/torchspec-2gpu-training}"
 model="${TORCHSPEC_CI_MODEL:-Qwen/Qwen3.8-27B}"
@@ -11,6 +12,9 @@ model_revision="${TORCHSPEC_CI_MODEL_REVISION:-1d4bf0f2ff6012fd82039f2fa52739d0d
 long_context_dataset_id="${TORCHSPEC_CI_LONG_CONTEXT_DATASET_ID:-long_chunked_prefill_test}"
 model_cache="${TORCHSPEC_CI_MODEL_CACHE:-${HF_HOME:-${artifact_dir}/huggingface}}"
 compile_cache="${TORCHSPEC_CI_COMPILE_CACHE:-${RUNNER_TEMP:-/tmp}/torchspec-torchinductor}"
+profile_num_steps="${TORCHSPEC_CI_PROFILE_NUM_STEPS:-}"
+profile_step_start="${TORCHSPEC_CI_PROFILE_STEP_START:-}"
+profile_step_end="${TORCHSPEC_CI_PROFILE_STEP_END:-}"
 
 mkdir -p "${artifact_dir}" "${model_cache}" "${compile_cache}"
 export HF_HOME="${model_cache}"
@@ -48,7 +52,7 @@ print(f"CI_SAMPLE_INPUT={sample['input']}")
 print(f"CI_SAMPLE_TARGET_OUTPUT={sample['target_output']}")
 PY
 
-expected_steps="$(python3 - "${fixture}" "${long_context_dataset_id}" <<'PY'
+fixture_records="$(python3 - "${fixture}" "${long_context_dataset_id}" <<'PY'
 import json
 import sys
 from pathlib import Path
@@ -70,7 +74,7 @@ print(len(records))
 PY
 )"
 echo "CI_LONG_CONTEXT_DATASET_ID=${long_context_dataset_id}"
-echo "CI_EPOCH_OPTIMIZER_STEPS=${expected_steps}"
+echo "CI_FIXTURE_RECORDS=${fixture_records}"
 
 if [[ -n "${TORCHSPEC_CI_MODEL_PATH:-}" ]]; then
   model_snapshot="${TORCHSPEC_CI_MODEL_PATH}"
@@ -103,22 +107,68 @@ echo "CI_MODEL_SNAPSHOT=${model_snapshot}"
 run_lane() {
   local lane="$1"
   local config="$2"
+  local expected_trainer="${3:-}"
   local lane_dir="${artifact_dir}/${lane}"
   local train_log="${lane_dir}/training.log"
+  local -a extra_args=()
 
   mkdir -p "${lane_dir}/actor-logs"
   export TORCHSPEC_LOG_DIR="${lane_dir}/actor-logs"
 
+  if [[ -n "${profile_num_steps}" ]]; then
+    [[ "${profile_num_steps}" =~ ^[1-9][0-9]*$ ]] || {
+      echo "TORCHSPEC_CI_PROFILE_NUM_STEPS must be a positive integer" >&2
+      exit 2
+    }
+    extra_args+=("training.num_train_steps=${profile_num_steps}")
+  fi
+  if [[ -n "${profile_step_start}" || -n "${profile_step_end}" ]]; then
+    [[ "${profile_step_start}" =~ ^[0-9]+$ && "${profile_step_end}" =~ ^[1-9][0-9]*$ ]] || {
+      echo "Both profile step bounds must be non-negative integers" >&2
+      exit 2
+    }
+    ((profile_step_end > profile_step_start)) || {
+      echo "Profile step end must be greater than profile step start" >&2
+      exit 2
+    }
+    mkdir -p "${lane_dir}/profiles"
+    extra_args+=(
+      "debug.enable_perf_metrics=true"
+      "debug.use_pytorch_profiler=true"
+      "debug.profile_target=[train_overall]"
+      "debug.profile_step_start=${profile_step_start}"
+      "debug.profile_step_end=${profile_step_end}"
+      "debug.profile_dir_name=${lane_dir}/profiles"
+    )
+  fi
+
   echo "CI_LANE_START=${lane}"
   cd "${repo_root}"
+  if [[ -n "${expected_trainer}" ]]; then
+    python3 - "${config}" "${expected_trainer}" <<'PY'
+import sys
+
+from torchspec import AutoDraftModelConfig
+from torchspec.config import load_config
+from torchspec.training.trainer_actor import _trainer_class_for_config
+
+config = load_config(sys.argv[1])
+draft_config = AutoDraftModelConfig.from_file(config.model.draft_model_config)
+trainer_name = _trainer_class_for_config(draft_config).__name__
+if trainer_name != sys.argv[2]:
+    raise SystemExit(f"Expected trainer {sys.argv[2]}, got {trainer_name}")
+print(f"CI_TRAINER lane={trainer_name.removesuffix('Trainer').lower()} class={trainer_name}")
+PY
+  fi
   python3 -m torchspec.train_entry \
     --config "${config}" \
     model.target_model_path="${model_snapshot}" \
     model_download_dir="${model_cache}" \
     cache_dir="${lane_dir}/cache" \
+    "${extra_args[@]}" \
     2>&1 | tee "${train_log}"
 
-  python3 - "${lane}" "${train_log}" "${lane_dir}/step-losses.json" "${expected_steps}" <<'PY'
+  python3 - "${lane}" "${train_log}" "${lane_dir}/step-losses.json" <<'PY'
 import json
 import math
 import re
@@ -126,20 +176,30 @@ import sys
 from pathlib import Path
 
 lane = sys.argv[1]
-pattern = re.compile(r"TRAIN_STEP step=(\d+) loss=([^ ]+)")
+lines = Path(sys.argv[2]).read_text(encoding="utf-8", errors="replace").splitlines()
+step_count_pattern = re.compile(r"num_train_steps=(\d+)")
+step_counts = [int(match.group(1)) for line in lines if (match := step_count_pattern.search(line))]
+if not step_counts:
+    raise SystemExit(f"{lane}: training log did not report num_train_steps")
+expected_step_count = step_counts[-1]
+if expected_step_count < 1:
+    raise SystemExit(f"{lane}: expected at least one optimizer step, got {expected_step_count}")
+
+loss_pattern = re.compile(r"TRAIN_STEP step=(\d+) loss=([^ ]+)")
 losses = []
-for line in Path(sys.argv[2]).read_text(encoding="utf-8", errors="replace").splitlines():
-    match = pattern.search(line)
+for line in lines:
+    match = loss_pattern.search(line)
     if match:
         losses.append({"step": int(match.group(1)), "loss": float(match.group(2))})
 
-expected_steps = list(range(1, int(sys.argv[4]) + 1))
+expected_steps = list(range(1, expected_step_count + 1))
 if [item["step"] for item in losses] != expected_steps:
     raise SystemExit(f"{lane}: expected losses for steps {expected_steps}, got {losses}")
 if not all(math.isfinite(item["loss"]) and item["loss"] > 0 for item in losses):
     raise SystemExit(f"{lane}: losses must be finite and positive: {losses}")
 
 Path(sys.argv[3]).write_text(json.dumps(losses, indent=2) + "\n", encoding="utf-8")
+print(f"CI_EPOCH_OPTIMIZER_STEPS lane={lane} count={expected_step_count}")
 print(f"CI_STEP_LOSSES lane={lane} values={json.dumps(losses, separators=(',', ':'))}")
 PY
   echo "CI_LANE_COMPLETE=${lane}"
@@ -147,3 +207,4 @@ PY
 
 run_lane eagle3 "${eagle3_config}"
 run_lane dspark "${dspark_config}"
+run_lane dflash2 "${dflash2_config}" DFlash2Trainer
diff --git a/torchspec/utils/env.py b/torchspec/utils/env.py
index b3a84191..fc60679a 100644
--- a/torchspec/utils/env.py
+++ b/torchspec/utils/env.py
@@ -23,8 +23,15 @@
     "TP_SOCKET_IFNAME",
     "CUTE_DSL_CACHE_DIR",
     "TORCHSPEC_FLASH_ATTN_OPT_LEVEL",
+    "VLLM_USE_V2_MODEL_RUNNER",
 ]
 
+_TORCHSPEC_DEFAULT_ENV_VARS = {
+    # TorchSpec's vLLM hidden-state extraction path requires Model Runner V2.
+    # An explicit process environment value still takes precedence below.
+    "VLLM_USE_V2_MODEL_RUNNER": "1",
+}
+
 # Prevent Ray from overriding VISIBLE_DEVICES so actors manage GPU assignment themselves.
 # Reference: https://github.com/ray-project/ray/blob/161849364/python/ray/_private/accelerators/
 _RAY_NOSET_VISIBLE_DEVICES_KEYS = [
@@ -49,6 +56,9 @@ def get_torchspec_env_vars() -> dict[str, str]:
     Intended for use with ``ray.remote(runtime_env={"env_vars": ...})``.
     Call-site env vars merged after this dict take higher priority.
     """
-    env = {k: "1" for k in _RAY_NOSET_VISIBLE_DEVICES_KEYS}
+    env = {
+        **_TORCHSPEC_DEFAULT_ENV_VARS,
+        **{k: "1" for k in _RAY_NOSET_VISIBLE_DEVICES_KEYS},
+    }
     env.update({k: os.environ[k] for k in _TORCHSPEC_ENV_KEYS if k in os.environ})
     return env